Skip to main content

AimStack

Aim is an open-source experiment tracking tool that enables users to log, visualize, and compare ML experiments.

Setting Up Aim in CGC​

Aim needs to be set up in the volume where you plan to train, evaluate and test your model. You can find information on how to create a volume here.

To set up Aim on your desired volume, provide the volume name in -v parameter and run the following command:

cgc compute create -n aim -c 2 -m 4 -v <train-vol-name> aimstack

This command will create a .aim repository at provided volume, where the tracked data will be stored. The Aim UI will be available at:

https://aim.<namespace>.cgc-waw-01.comtegra.cloud/aim

Integrate Aim into your project​

To use Aim in your project, you first need to install it using pip:

pip install aim

Below is an example of how to integrate Aim into a YOLOv8 training pipeline to track metrics and hyperparameters effectively.

from ultralytics import YOLO
import aim

# Initialize an Aim run with a custom experiment name to group related runs together
run = aim.Run(
    experiment='experiment1'
)

# Set training hyperparameters
run["hyperparameters"] = {
    "epochs": 2,
    "batch_size": 32,
    "learning_rate": 0.001429,
    "img_size": 720
}

# Load the YOLOv8 model
model = YOLO('yolov8n.pt')

def on_train_epoch_end(trainer):
    epoch = trainer.epoch
    metrics = trainer.metrics
    
    if metrics:
        run.track(metrics['val/box_loss'], name="Validation Box Loss", epoch=epoch)
        run.track(metrics['val/cls_loss'], name="Validation Classification Loss", epoch=epoch)
        run.track(metrics['metrics/precision(B)'], name="Precision", epoch=epoch)
        run.track(metrics['metrics/recall(B)'], name="Recall", epoch=epoch)
        run.track(metrics['metrics/mAP50(B)'], name="mAP@0.5", epoch=epoch)
    else:
        print(f"No metrics available at epoch {epoch}")

# Add Aim tracking as a callback to the model
model.add_callback('on_train_epoch_end', on_train_epoch_end)

# Start training and log parameters
model.train(
    data='/workspace/bicycle-tracks/data/data.yaml',
    epochs=run['hyperparameters']["epochs"],
    imgsz=run['hyperparameters']["img_size"],
    batch=run['hyperparameters']["batch_size"],
    workers=2,
    verbose=False
)

# Close the Aim run
run.close()

You can see the results in Aim app you've created earlier. In this example the results will be visible at the end of each epoch.

To learn more about Aim and how to use it effectively, refer to the following links:

These resources provide in-depth explanations of features, best practices, and additional integrations for experiment tracking with Aim.