View source Download .ipynb

Train Yolo-NAS model with 3LC Metrics Collection

This notebook shows how to train a Yolo-NAS detection model with 3LC Metrics Collection on an Object Detection 3LC Table.

We recommend to use Python 3.10, and ensure super-gradients is installed. Note that installing super-gradients will downgrade the versions of several dependencies of 3lc. We therefore install termcolor==3.1.0 afterwards.

Install dependencies

%pip install 3lc
%pip install super-gradients
%pip install termcolor==3.1.0

Imports

import tlc
from super_gradients.common.object_names import Models
from super_gradients.training import Trainer, models
from super_gradients.training.losses import PPYoloELoss
from super_gradients.training.metrics import DetectionMetrics_050
from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback
from super_gradients.training.transforms.transforms import (
    DetectionHorizontalFlip,
    DetectionHSV,
    DetectionPaddedRescale,
    DetectionRandomAffine,
    DetectionStandardize,
    DetectionTargetsFormatTransform,
)
from super_gradients.training.utils.collate_fn.detection_collate_fn import DetectionCollateFN
from tlc.integration.super_gradients import (
    DetectionDataset,
    DetectionMetricsCollectionCallback,
    PipelineParams,
)
from torch.utils.data import DataLoader

Project Setup

PROJECT_NAME = "3LC Tutorials"
DATASET_NAME = "COCO128"

We reuse the COCO128 tables created in create-table-from-coco-detection.ipynb, so make sure to run this notebook first.

train_table = tlc.Table.from_names("initial", DATASET_NAME, PROJECT_NAME)
val_table = tlc.Table.from_names("initial", DATASET_NAME, PROJECT_NAME)

simple_value_map = train_table.get_simple_value_map("bbs.bb_list.label")
class_names = list(simple_value_map.values())
num_classes = len(class_names)
class_names

Pass the training and validation tables to a DetectionDataset, which loads images and annotations from the Table and retains image and label preprocessing functionality available in the parent SuperGradients DetectionDataset.

train_dataset = DetectionDataset(
    train_table,
    input_dim=(640, 640),
    transforms=[
        DetectionRandomAffine(
            degrees=0.0,
            scales=(0.5, 1.5),
            shear=0.0,
            target_size=(640, 640),
            filter_box_candidates=False,
            border_value=128,
        ),
        DetectionHSV(prob=1.0, hgain=5, vgain=30, sgain=30),
        DetectionHorizontalFlip(prob=0.5),
        DetectionPaddedRescale(input_dim=(640, 640)),
        DetectionStandardize(max_value=255),
        DetectionTargetsFormatTransform(input_dim=(640, 640), output_format="LABEL_CXCYWH"),
    ],
)
val_dataset = DetectionDataset(
    val_table,
    input_dim=(640, 640),
    transforms=[
        DetectionPaddedRescale(input_dim=(640, 640), max_targets=300),
        DetectionStandardize(max_value=255),
        DetectionTargetsFormatTransform(input_dim=(640, 640), output_format="LABEL_CXCYWH"),
    ],
)

train_dataloader: DataLoader = DataLoader(train_dataset, batch_size=4, shuffle=True, collate_fn=DetectionCollateFN())
val_dataloader: DataLoader = DataLoader(val_dataset, batch_size=4, shuffle=False, collate_fn=DetectionCollateFN())

model = models.get(Models.YOLO_NAS_S, num_classes=num_classes, pretrained_weights="coco")

trainer = Trainer(
    experiment_name="fine-tune-yolo-nas-detect-3lc",
    ckpt_root_dir="checkpoints",
)

Use PipelineParams to customize the SuperGradients Pipeline used for inference when collecting metrics, and a DetectionMetricsCollectionCallback to invoke 3LC Metrics Collection, Hyperparameter and Run logging. Make sure to pass the callback to the phase_callbacks in the training parameters passed to the SuperGradients Trainer.

pipeline_params = PipelineParams(
    iou=0.5,  # IoU threshold for NMS matching
    conf=0.3,  # Confidence threshold for predictions to keep
    max_predictions=10,  # Max number of predictions for each image
    class_agnostic_nms=True,
)

callback = DetectionMetricsCollectionCallback(
    project_name=PROJECT_NAME,
    run_description="Fine-tune Yolo-NAS detection model",
    collect_val_only=True,
    pipeline_params=pipeline_params,
)

train_params = {
    "warmup_initial_lr": 5e-5,
    "initial_lr": 5e-5,
    "lr_mode": "cosine",
    "cosine_final_lr_ratio": 0.5,
    "optimizer": "AdamW",
    "zero_weight_decay_on_bias_and_bn": True,
    "lr_warmup_epochs": 1,
    "warmup_mode": "LinearEpochLRWarmup",
    "optimizer_params": {"weight_decay": 0.0001},
    "ema": False,
    "average_best_models": False,
    "ema_params": {"beta": 25, "decay_type": "exp"},
    "max_epochs": 5,
    "mixed_precision": True,
    "loss": PPYoloELoss(use_static_assigner=False, num_classes=num_classes, reg_max=16),
    "valid_metrics_list": [
        DetectionMetrics_050(
            score_thres=0.1,
            top_k_predictions=300,
            num_cls=num_classes,
            normalize_targets=True,
            include_classwise_ap=True,
            class_names=class_names,
            post_prediction_callback=PPYoloEPostPredictionCallback(
                score_threshold=0.01, nms_top_k=1000, max_predictions=300, nms_threshold=0.7
            ),
        )
    ],
    "metric_to_watch": "mAP@0.50",
    "phase_callbacks": [callback],
}

run = tlc.init(project_name=PROJECT_NAME)

trainer.train(
    model=model,
    training_params=train_params,
    train_loader=train_dataloader,
    valid_loader=val_dataloader,
)