基于PyTorch的YOLO NAS-s模型训练后无检测框问题排查

我用PyTorch训练了一款针对零售门店图像检测的YOLO NAS-s模型，数据集含约8000张图片，计划训练20个epoch。训练因内存问题中断后，从checkpoint文件恢复训练，但测试时仅加载图片，没有生成任何检测框。

数据集路径声明代码 #

from super_gradients.training.dataloaders.dataloaders import coco_detection_yolo_format_train, coco_detection_yolo_format_val
    
BATCH_SIZE = 1
CLASSES = ['product']
CLASSES += [str(i) for i in range(80 - len(CLASSES))]
    
dataset_params = {
    'data_dir': r"C:\Users\Giriraj\Documents\Prernas ML Models\SKU110KDataset\SKU110K_fixed",
    'train_images_dir':r'C:\Users\Giriraj\Documents\Prernas ML Models\SKU110KDataset\SKU110K_fixed\images\train',
    'train_labels_dir':r'C:\Users\Giriraj\Documents\Prernas ML Models\SKU110KDataset\SKU110K_fixed\labels\train',
    'val_images_dir':r'C:\Users\Giriraj\Documents\Prernas ML Models\SKU110KDataset\SKU110K_fixed\images\val',
    'val_labels_dir':r'C:\Users\Giriraj\Documents\Prernas ML Models\SKU110KDataset\SKU110K_fixed\labels\val',
    'test_images_dir':r'C:\Users\Giriraj\Documents\Prernas ML Models\SKU110KDataset\SKU110K_fixed\images\test',
    'test_labels_dir':r'C:\Users\Giriraj\Documents\Prernas ML Models\SKU110KDataset\SKU110K_fixed\labels\test',
    'classes': CLASSES
}
    
train_data = coco_detection_yolo_format_train(
    dataset_params={
        'data_dir': dataset_params['data_dir'],
        'images_dir': dataset_params['train_images_dir'],
        'labels_dir': dataset_params['train_labels_dir'],
        'classes': dataset_params['classes']
    },
    dataloader_params={
        'batch_size': BATCH_SIZE,
        'num_workers': 2
    }
)
    
val_data = coco_detection_yolo_format_val(
    dataset_params={
        'data_dir': dataset_params['data_dir'],
        'images_dir': dataset_params['val_images_dir'],
        'labels_dir': dataset_params['val_labels_dir'],
        'classes': dataset_params['classes']
    },
    dataloader_params={
        'batch_size': BATCH_SIZE,
        'num_workers': 2
    }
)
    
test_data = coco_detection_yolo_format_val(
    dataset_params={
        'data_dir': dataset_params['data_dir'],
        'images_dir': dataset_params['test_images_dir'],
        'labels_dir': dataset_params['test_labels_dir'],
        'classes': dataset_params['classes']
    },
    dataloader_params={
        'batch_size': BATCH_SIZE,
        'num_workers': 2
    }
)

模型与设备声明代码 #

import torch
from super_gradients.training import models
from super_gradients.training import Trainer

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
model = models.get('yolo_nas_s', pretrained_weights="coco").to(DEVICE)
#optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
#optimizer = torch.optim.Adam(model.parameters())
trainer = Trainer(experiment_name="SKU110K", ckpt_root_dir="./weights")

训练参数配置代码 #

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
    
MAX_EPOCHS = 20
    
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,
    "resume":True,
    "ema_params": {"decay": 0.9, "decay_type": "threshold"},
    "max_epochs": MAX_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=50,
            num_cls=len(dataset_params['classes']),
            normalize_targets=True,
            post_prediction_callback=PPYoloEPostPredictionCallback(
                score_threshold=0.01,
                nms_top_k=100,
                max_predictions=20,
                nms_threshold=0.7
            )
        )
    ],
    "metric_to_watch": 'mAP@0.50'
}

训练与恢复训练代码 #

import torch
    
# Set the desired starting epoch
    
# Load the model with the weights from the average checkpoint file
checkpoint_file = torch.load(r'C:\Users\Giriraj\Documents\Prernas ML Models\yolo-nas-retail-training-main\yolo-nas-retail-training-main\weights\SKU110K\average_model.pth')
model.load_state_dict(checkpoint_file['net'])
START_EPOCH = checkpoint_file['epoch'] + 1
max_epochs=20
    
# Your training loop
for epoch in range(START_EPOCH, max_epochs):
    # Train the model
    trainer.train(
        model=model,
        training_params=train_params, 
        train_loader=train_data, 
        valid_loader=val_data
    )       

可能原因及解决办法 #

1. 模型未切换到评估模式 #

测试时模型若处于训练模式，Batch Normalization、Dropout等层的行为会干扰检测结果，必须先执行：

model.eval()

2. Checkpoint加载不完整 #

训练参数中开启了ema=True，checkpoint内包含EMA权重（性能更优），但你只加载了net权重。可改用Trainer内置方法加载：

trainer.load_checkpoint(checkpoint_path='./weights/SKU110K/average_model.pth', load_best=True)

或手动加载EMA权重：

if 'ema_net' in checkpoint_file:
    model.load_state_dict(checkpoint_file['ema_net'])

3. 恢复训练逻辑错误 #

当前循环每次调用trainer.train()都会从epoch 0重新开始，而非从START_EPOCH延续。正确做法是在train_params中指定恢复路径，无需手动写循环：

train_params['resume'] = True
train_params['resume_path'] = './weights/SKU110K/average_model.pth'

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

4. 测试后处理参数设置问题 #

训练时后处理的score_threshold=0.01，若测试时阈值过高会过滤所有结果。测试时需保持参数一致，或临时降低阈值排查：

from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback

# 手动指定后处理参数
post_process = PPYoloEPostPredictionCallback(
    score_threshold=0.01,
    nms_top_k=100,
    max_predictions=20,
    nms_threshold=0.7
)

# 推理后应用后处理
predictions = model.predict(test_image_path)
processed_predictions = post_process(predictions)

5. 数据集类别配置冗余 #

SKU110K仅1类，但你给CLASSES补充了79个冗余类别，会导致模型输出混乱，应改为：

CLASSES = ['product']

内容的提问来源于stack exchange，提问作者Prernq