YOLO-NAS 自定义目标检测训练实战：SuperGradients 训练与测试-尧图网络科技

YOLO-NAS 自定义目标检测训练实战：SuperGradients 训练与测试

这篇教程根据我复现 YOLO-NAS 自定义目标检测训练流程时整理，重点演示如何安装 SuperGradients、加载预训练模型、训练自定义数据集并做测试集评估。

YOLO-NAS 的训练方式和 Ultralytics 不同，依赖 SuperGradients 的训练器和数据加载器。本文适合想把 YOLO-NAS 接入自定义训练流程的同学。

本文会重点跑通以下流程：

安装 SuperGradients 与可视化依赖
下载示例图片并跑通预训练推理
从数据集后台获取 YOLO 格式数据集
配置训练参数并启动训练
用测试集和混淆矩阵检查模型表现

如果你正在系统学习目标检测、实例分割、OCR、多目标跟踪或视觉大模型，建议收藏本文；配套 notebook、示例图片和运行环境说明后续会继续整理。如果环境配置卡住，可以在评论区说明具体报错。

📚 文章目录

YOLO-NAS 自定义目标检测训练实战：SuperGradients 训练与测试
- ⚙️ 环境准备
- ⬇️ 下载权重与示例图片
- 🔍 预训练模型推理
- 📦 从数据集后台获取 YOLO 数据集
- 🧱 配置训练参数
- 🏋️ 启动训练
- 📏 训练后测试
- 🧪 批量预测可视化
- 📊 结果评估
- 📌 小结
- 📚 同系列教程汇总

⚙️ 环境准备

先检查运行环境并安装依赖。建议优先使用带 NVIDIA GPU 的环境，避免推理和训练阶段显存不足。

!nvidia-smi

!pip install-q supervision==0.25.0# URLs changed for YOLO-NAS models, but the authors are not updating the repo# Issue: https://github.com/Deci-AI/super-gradients/issues/2057# Fix PR: https://github.com/Deci-AI/super-gradients/pull/2061# Option 1 - Install from patched branch. Does not work, causes different errors!# !pip install -q git+https://github.com/hannadiamond/super-gradients@patch-1# Option 2 - Install normally, and modify code# Pip errors are expected.!pip install-q git+https://github.com/Deci-AI/super-gradients.git@stable !sed-i-e"s/sghub.deci.ai/sg-hub-nv.s3.amazonaws.com/g"/usr/local/lib/python3.10/dist-packages/super_gradients/training/pretrained_models.py !sed-i-e"s/sghub.deci.ai/sg-hub-nv.s3.amazonaws.com/g"/usr/local/lib/python3.10/dist-packages/super_gradients/training/utils/checkpoint_utils.py

# URLs changed for YOLO-NAS models, but the authors aren't updating the repo# Issue: https://github.com/Deci-AI/super-gradients/issues/2057# Fix PR: https://github.com/Deci-AI/super-gradients/pull/2061# Option 1 - Install from patched branch. Does not work, causes different errors!# !pip install -q git+https://github.com/hannadiamond/super-gradients@patch-1# Option 2 - Install normally, and modify code# Pip errors are expected.!pip install-q git+https://github.com/Deci-AI/super-gradients.git@stable !sed-i-e"s/sghub.deci.ai/sg-hub-nv.s3.amazonaws.com/g"/usr/local/lib/python3.10/dist-packages/super_gradients/training/pretrained_models.py !sed-i-e"s/sghub.deci.ai/sg-hub-nv.s3.amazonaws.com/g"/usr/local/lib/python3.10/dist-packages/super_gradients/training/utils/checkpoint_utils.py

importos HOME=os.getcwd()print(HOME)

⬇️ 下载权重与示例图片

先把 SAM checkpoint 和示例图片准备好。

importtorch DEVICE='cuda'iftorch.cuda.is_available()else"cpu"MODEL_ARCH='yolo_nas_l'

fromsuper_gradients.trainingimportmodels model=models.get(MODEL_ARCH,pretrained_weights="coco").to(DEVICE)

f"{HOME}/data"

%cd{HOME}!mkdir{HOME}/data%cd{HOME}/data# 请从数据集后台下载示例图片，保存为 dog.jpeg 到 dog-8.jpeg。

🔍 预训练模型推理

用预训练模型在单张图上跑一遍检测结果。

SOURCE_IMAGE_PATH=f"{HOME}/data/dog-3.jpeg"

importcv2 image=cv2.imread(SOURCE_IMAGE_PATH)result=model.predict(image,conf=0.35)

type(result)

importsupervisionassv detections=sv.Detections(xyxy=result.prediction.bboxes_xyxy,confidence=result.prediction.confidence,class_id=result.prediction.labels.astype(int))box_annotator=sv.BoxAnnotator()label_annotator=sv.LabelAnnotator()labels=[f"{result.class_names[class_id]}{confidence:0.2f}"for_,_,confidence,class_id,_,_indetections]annotated_frame=image.copy()annotated_frame=box_annotator.annotate(scene=annotated_frame,detections=detections)annotated_frame=label_annotator.annotate(scene=annotated_frame,detections=detections,labels=labels)%matplotlib inline sv.plot_image(annotated_frame,(12,12))

📦 从数据集后台获取 YOLO 数据集

从数据集后台导出 YOLO 格式数据集后，训练时直接引用本地路径。

%cd{HOME}# 如需使用自定义数据集，请从数据集后台下载 YOLO 格式数据并解压到本地目录。fromtypesimportSimpleNamespace DATASET_DIR="/content/dataset"# 修改为数据集后台导出的数据集目录dataset=SimpleNamespace(location=DATASET_DIR)

LOCATION=dataset.locationprint("location:",LOCATION)CLASSES=sorted(project.classes.keys())print("classes:",CLASSES)

🧱 配置训练参数

把 batch、epoch、类别和路径整理成训练参数。

MODEL_ARCH='yolo_nas_l'BATCH_SIZE=8MAX_EPOCHS=25CHECKPOINT_DIR=f'{HOME}/checkpoints'EXPERIMENT_NAME=project.name.lower().replace(" ","_")

fromsuper_gradients.trainingimportTrainer trainer=Trainer(experiment_name=EXPERIMENT_NAME,ckpt_root_dir=CHECKPOINT_DIR)

dataset_params={'data_dir':LOCATION,'train_images_dir':'train/images','train_labels_dir':'train/labels','val_images_dir':'valid/images','val_labels_dir':'valid/labels','test_images_dir':'test/images','test_labels_dir':'test/labels','classes':CLASSES}

fromsuper_gradients.training.dataloaders.dataloadersimport(coco_detection_yolo_format_train,coco_detection_yolo_format_val)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})

train_data.dataset.transforms

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

fromsuper_gradients.training.lossesimportPPYoloELossfromsuper_gradients.training.metricsimportDetectionMetrics_050fromsuper_gradients.training.models.detection_models.pp_yolo_eimportPPYoloEPostPredictionCallback 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":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=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'}

🏋️ 启动训练

训练器和数据加载器准备好后，就可以直接开始训练。

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

%load_ext tensorboard%tensorboard--logdir{CHECKPOINT_DIR}/{EXPERIMENT_NAME}

!zip-r yolo_nas.zip{CHECKPOINT_DIR}/{EXPERIMENT_NAME}

# if you experience 'NotImplementedError: A UTF-8 locale is required. Got ANSI_X3.4-1968' error, run code below 👇# import locale# locale.getpreferredencoding = lambda: "UTF-8"

📏 训练后测试

用测试集检查 average model 或 best model 的表现。

best_model=models.get(MODEL_ARCH,num_classes=len(dataset_params['classes']),checkpoint_path=f"{CHECKPOINT_DIR}/{EXPERIMENT_NAME}/average_model.pth").to(DEVICE)

trainer.test(model=best_model,test_loader=test_data,test_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)))

🧪 批量预测可视化

把测试集中的样本和预测结果并排显示，方便人工检查。

importsupervisionassv ds=sv.DetectionDataset.from_yolo(images_directory_path=f"{dataset.location}/test/images",annotations_directory_path=f"{dataset.location}/test/labels",data_yaml_path=f"{dataset.location}/data.yaml",force_masks=False)

importsupervisionassv CONFIDENCE_TRESHOLD=0.5predictions={}forimage_name,imageinds.images.items():result=list(best_model.predict(image,conf=CONFIDENCE_TRESHOLD))[0]detections=sv.Detections(xyxy=result.prediction.bboxes_xyxy,confidence=result.prediction.confidence,class_id=result.prediction.labels.astype(int))predictions[image_name]=detections

importrandom random.seed(10)

importsupervisionassv MAX_IMAGE_COUNT=5n=min(MAX_IMAGE_COUNT,len(ds.images))keys=list(ds.images.keys())keys=random.sample(keys,n)box_annotator=sv.BoxAnnotator()images=[]titles=[]forkeyinkeys:frame_with_annotations=box_annotator.annotate(scene=ds.images[key].copy(),detections=ds.annotations[key],skip_label=True)images.append(frame_with_annotations)titles.append('annotations')frame_with_predictions=box_annotator.annotate(scene=ds.images[key].copy(),detections=predictions[key],skip_label=True)images.append(frame_with_predictions)titles.append('predictions')%matplotlib inline sv.plot_images_grid(images=images,titles=titles,grid_size=(n,2),size=(2*4,n*4))

📊 结果评估

最后补上混淆矩阵，整体看一遍分类和检测效果。

!pip install onemetric

importosimportnumpyasnpfromonemetric.cv.object_detectionimportConfusionMatrix keys=list(ds.images.keys())annotation_batches,prediction_batches=[],[]forkeyinkeys:annotation=ds.annotations[key]annotation_batch=np.column_stack((annotation.xyxy,annotation.class_id))annotation_batches.append(annotation_batch)prediction=predictions[key]prediction_batch=np.column_stack((prediction.xyxy,prediction.class_id,prediction.confidence))prediction_batches.append(prediction_batch)confusion_matrix=ConfusionMatrix.from_detections(true_batches=annotation_batches,detection_batches=prediction_batches,num_classes=len(ds.classes),conf_threshold=CONFIDENCE_TRESHOLD)confusion_matrix.plot(os.path.join(HOME,"confusion_matrix.png"),class_names=ds.classes)