【视觉算法系列3】在自定义数据集上训练 YOLO NAS（下篇）

在这里插入图片描述

提示：免费获取本文涉及的完整代码与数据集，请添加微信peaeci122

YOLO-NAS是目前最新的YOLO目标检测模型，它在准确性方面击败了所有其他 YOLO 模型。与之前的 YOLO 模型相比，预训练的 YOLO-NAS 模型能够以更高的准确度检测更多目标。

如何在自定义数据集上训练 YOLO NAS？这将是我这两篇文章的目标。

1、用于训练 YOLO NAS 的物体检测数据集

2、在自定义数据集上训练 YOLO NAS

3、微调 YOLO NAS 模型

4、使用经过训练的 YOLO NAS 模型对测试图像进行推理

5、YOLO NAS 训练模型视频推理结果

6、结论

上篇讲到“微调 YOLO NAS 模型”中的“定义YOLO NAS训练的变换和增强”，本篇我们继续

YOLO NAS 训练参数、模型培训、分析微调结果。

YOLO NAS 训练参数

在我们开始微调过程之前，训练参数是最重要的组成部分。在这里，我们定义了要训练的epoch、要监控的验证指标和学习率等。

train_params = {
    'silent_mode': False,
    "average_best_models":True,
    "warmup_mode": "linear_epoch_step",
    "warmup_initial_lr": 1e-6,
    "lr_warmup_epochs": 3,
    "initial_lr": 5e-4,
    "lr_mode": "cosine",
    "cosine_final_lr_ratio": 0.1,
    "optimizer": "Adam",
    "optimizer_params": {"weight_decay": 0.0001},
    "zero_weight_decay_on_bias_and_bn": True,
    "ema": True,
    "ema_params": {"decay": 0.9, "decay_type": "threshold"},
    "max_epochs": EPOCHS,
    "mixed_precision": True,
    "loss": PPYoloELoss(
        use_static_assigner=False,
        num_classes=len(dataset_params['classes']),
        reg_max=16
    ),
    "valid_metrics_list": [
        DetectionMetrics_050(
            score_thres=0.1,
            top_k_predictions=300,
            num_cls=len(dataset_params['classes']),
            normalize_targets=True,
            post_prediction_callback=PPYoloEPostPredictionCallback(
                score_threshold=0.01,
                nms_top_k=1000,
                max_predictions=300,
                nms_threshold=0.7
            )
        ),
        DetectionMetrics_050_095(
            score_thres=0.1,
            top_k_predictions=300,
            num_cls=len(dataset_params['classes']),
            normalize_targets=True,
            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:0.95'
}

在训练时，输出将同时显示 50% Iou 和 5%-95% IoU 时的 mAP。不过，我们只监控主要指标（mAP@0.50:0.95 IoU），因此最好的模型将根据它被保存。、

YOLO NAS 模型培训

由于我们要训练三个不同的模型，因此需要将过程自动化一些，可以定义一个包含三个模型名称的列表，并根据该列表设置检查点目录。这也会加载适当的模型，因为列表中的模型名称与super-gradients API中的模型名称匹配。

models_to_train = [
    'yolo_nas_s',
    'yolo_nas_m',
    'yolo_nas_l'
]

CHECKPOINT_DIR = 'checkpoints'

for model_to_train in models_to_train:
    trainer = Trainer(
        experiment_name=model_to_train, 
        ckpt_root_dir=CHECKPOINT_DIR
    )

    model = models.get(
        model_to_train, 
        num_classes=len(dataset_params['classes']), 
        pretrained_weights="coco"
    )

    trainer.train(
        model=model, 
        training_params=train_params, 
        train_loader=train_data, 
        valid_loader=val_data
    )

三个训练实验将依次运行，所有的模型检查点将保存在各自的目录中。

分析 YOLO NAS 微调结果

在训练过程中，输出单元 / 终端显示训练过程的全面视图。

SUMMARY OF EPOCH 0
├── Training
│   ├── Ppyoloeloss/loss = 3.8575
│   ├── Ppyoloeloss/loss_cls = 2.3712
│   ├── Ppyoloeloss/loss_dfl = 1.1773
│   └── Ppyoloeloss/loss_iou = 0.3591
└── Validation
    ├── F1@0.50 = 0.0
    ├── F1@0.50:0.95 = 0.0
    ├── Map@0.50 = 0.0012
    ├── Map@0.50:0.95 = 0.0005
    ├── Ppyoloeloss/loss = 3.7911
    ├── Ppyoloeloss/loss_cls = 2.5251
    ├── Ppyoloeloss/loss_dfl = 0.9791
    ├── Ppyoloeloss/loss_iou = 0.3106
    ├── Precision@0.50 = 0.0
    ├── Precision@0.50:0.95 = 0.0
    ├── Recall@0.50 = 0.0
    └── Recall@0.50:0.95 = 0.0
.
.
.
SUMMARY OF EPOCH 50
├── Training
│   ├── Ppyoloeloss/loss = 1.4382
│   │   ├── Best until now = 1.433  (↗ 0.0053)
│   │   └── Epoch N-1      = 1.433  (↗ 0.0053)
│   ├── Ppyoloeloss/loss_cls = 0.6696
│   │   ├── Best until now = 0.6651 (↗ 0.0046)
│   │   └── Epoch N-1      = 0.6651 (↗ 0.0046)
│   ├── Ppyoloeloss/loss_dfl = 0.6859
│   │   ├── Best until now = 0.6846 (↗ 0.0013)
│   │   └── Epoch N-1      = 0.686  (↘ -0.0)
│   └── Ppyoloeloss/loss_iou = 0.1703
│       ├── Best until now = 0.17   (↗ 0.0003)
│       └── Epoch N-1      = 0.17   (↗ 0.0003)
└── Validation
    ├── F1@0.50 = 0.292
    │   ├── Best until now = 0.3025 (↘ -0.0104)
    │   └── Epoch N-1      = 0.2774 (↗ 0.0146)
    ├── F1@0.50:0.95 = 0.1859
    │   ├── Best until now = 0.1928 (↘ -0.007)
    │   └── Epoch N-1      = 0.1761 (↗ 0.0097)
    ├── Map@0.50 = 0.7631
    │   ├── Best until now = 0.7745 (↘ -0.0114)
    │   └── Epoch N-1      = 0.7159 (↗ 0.0472)
    ├── Map@0.50:0.95 = 0.4411
    │   ├── Best until now = 0.4443 (↘ -0.0032)
    │   └── Epoch N-1      = 0.4146 (↗ 0.0265)
    ├── Ppyoloeloss/loss = 1.5389
    │   ├── Best until now = 1.5404 (↘ -0.0015)
    │   └── Epoch N-1      = 1.5526 (↘ -0.0137)
    ├── Ppyoloeloss/loss_cls = 0.6893
    │   ├── Best until now = 0.687  (↗ 0.0024)
    │   └── Epoch N-1      = 0.6972 (↘ -0.0079)
    ├── Ppyoloeloss/loss_dfl = 0.7148
    │   ├── Best until now = 0.7136 (↗ 0.0012)
    │   └── Epoch N-1      = 0.7234 (↘ -0.0086)
    ├── Ppyoloeloss/loss_iou = 0.1969
    │   ├── Best until now = 0.1953 (↗ 0.0016)
    │   └── Epoch N-1      = 0.1975 (↘ -0.0006)
    ├── Precision@0.50 = 0.1828
    │   ├── Best until now = 0.1926 (↘ -0.0097)
    │   └── Epoch N-1      = 0.1718 (↗ 0.011)
    ├── Precision@0.50:0.95 = 0.1166
    │   ├── Best until now = 0.1229 (↘ -0.0063)
    │   └── Epoch N-1      = 0.1092 (↗ 0.0074)
    ├── Recall@0.50 = 0.8159
    │   ├── Best until now = 0.8939 (↘ -0.0781)
    │   └── Epoch N-1      = 0.8307 (↘ -0.0149)
    └── Recall@0.50:0.95 = 0.522
        ├── Best until now = 0.5454 (↘ -0.0234)
        └── Epoch N-1      = 0.5236 (↘ -0.0016)

===========================================================

我们可以查看 Tensorboard 日志并检查 mAP 图，以便比较三个训练全部过程。Tensorboard 日志位于检查点目录下的相应训练文件夹中。

下图显示了三个训练实验的主要AP对比：

图 5.微调后的 YOLO NAS 小型、中型和大型模型对比