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 supervision0.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 githttps://github.com/hannadiamond/super-gradientspatch-1# Option 2 - Install normally, and modify code# Pip errors are expected.!pip install-q githttps://github.com/Deci-AI/super-gradients.gitstable !sed-i-es/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-es/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 arent 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 githttps://github.com/hannadiamond/super-gradientspatch-1# Option 2 - Install normally, and modify code# Pip errors are expected.!pip install-q githttps://github.com/Deci-AI/super-gradients.gitstable !sed-i-es/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-es/sghub.deci.ai/sg-hub-nv.s3.amazonaws.com/g/usr/local/lib/python3.10/dist-packages/super_gradients/training/utils/checkpoint_utils.pyimportos HOMEos.getcwd()print(HOME)⬇️ 下载权重与示例图片先把 SAM checkpoint 和示例图片准备好。importtorch DEVICEcudaiftorch.cuda.is_available()elsecpuMODEL_ARCHyolo_nas_lfromsuper_gradients.trainingimportmodels modelmodels.get(MODEL_ARCH,pretrained_weightscoco).to(DEVICE)f{HOME}/data%cd{HOME}!mkdir{HOME}/data%cd{HOME}/data# 请从数据集后台下载示例图片保存为 dog.jpeg 到 dog-8.jpeg。 预训练模型推理用预训练模型在单张图上跑一遍检测结果。SOURCE_IMAGE_PATHf{HOME}/data/dog-3.jpegimportcv2 imagecv2.imread(SOURCE_IMAGE_PATH)resultmodel.predict(image,conf0.35)type(result)importsupervisionassv detectionssv.Detections(xyxyresult.prediction.bboxes_xyxy,confidenceresult.prediction.confidence,class_idresult.prediction.labels.astype(int))box_annotatorsv.BoxAnnotator()label_annotatorsv.LabelAnnotator()labels[f{result.class_names[class_id]}{confidence:0.2f}for_,_,confidence,class_id,_,_indetections]annotated_frameimage.copy()annotated_framebox_annotator.annotate(sceneannotated_frame,detectionsdetections)annotated_framelabel_annotator.annotate(sceneannotated_frame,detectionsdetections,labelslabels)%matplotlib inline sv.plot_image(annotated_frame,(12,12)) 从数据集后台获取 YOLO 数据集从数据集后台导出 YOLO 格式数据集后训练时直接引用本地路径。%cd{HOME}# 如需使用自定义数据集请从数据集后台下载 YOLO 格式数据并解压到本地目录。fromtypesimportSimpleNamespace DATASET_DIR/content/dataset# 修改为数据集后台导出的数据集目录datasetSimpleNamespace(locationDATASET_DIR)LOCATIONdataset.locationprint(location:,LOCATION)CLASSESsorted(project.classes.keys())print(classes:,CLASSES) 配置训练参数把 batch、epoch、类别和路径整理成训练参数。MODEL_ARCHyolo_nas_lBATCH_SIZE8MAX_EPOCHS25CHECKPOINT_DIRf{HOME}/checkpointsEXPERIMENT_NAMEproject.name.lower().replace( ,_)fromsuper_gradients.trainingimportTrainer trainerTrainer(experiment_nameEXPERIMENT_NAME,ckpt_root_dirCHECKPOINT_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_datacoco_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_datacoco_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_datacoco_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.transformsfromsuper_gradients.trainingimportmodels modelmodels.get(MODEL_ARCH,num_classeslen(dataset_params[classes]),pretrained_weightscoco)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_assignerFalse,num_classeslen(dataset_params[classes]),reg_max16),valid_metrics_list:[DetectionMetrics_050(score_thres0.1,top_k_predictions300,num_clslen(dataset_params[classes]),normalize_targetsTrue,post_prediction_callbackPPYoloEPostPredictionCallback(score_threshold0.01,nms_top_k1000,max_predictions300,nms_threshold0.7))],metric_to_watch:mAP0.50}️ 启动训练训练器和数据加载器准备好后就可以直接开始训练。trainer.train(modelmodel,training_paramstrain_params,train_loadertrain_data,valid_loaderval_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_modelmodels.get(MODEL_ARCH,num_classeslen(dataset_params[classes]),checkpoint_pathf{CHECKPOINT_DIR}/{EXPERIMENT_NAME}/average_model.pth).to(DEVICE)trainer.test(modelbest_model,test_loadertest_data,test_metrics_listDetectionMetrics_050(score_thres0.1,top_k_predictions300,num_clslen(dataset_params[classes]),normalize_targetsTrue,post_prediction_callbackPPYoloEPostPredictionCallback(score_threshold0.01,nms_top_k1000,max_predictions300,nms_threshold0.7))) 批量预测可视化把测试集中的样本和预测结果并排显示方便人工检查。importsupervisionassv dssv.DetectionDataset.from_yolo(images_directory_pathf{dataset.location}/test/images,annotations_directory_pathf{dataset.location}/test/labels,data_yaml_pathf{dataset.location}/data.yaml,force_masksFalse)importsupervisionassv CONFIDENCE_TRESHOLD0.5predictions{}forimage_name,imageinds.images.items():resultlist(best_model.predict(image,confCONFIDENCE_TRESHOLD))[0]detectionssv.Detections(xyxyresult.prediction.bboxes_xyxy,confidenceresult.prediction.confidence,class_idresult.prediction.labels.astype(int))predictions[image_name]detectionsimportrandom random.seed(10)importsupervisionassv MAX_IMAGE_COUNT5nmin(MAX_IMAGE_COUNT,len(ds.images))keyslist(ds.images.keys())keysrandom.sample(keys,n)box_annotatorsv.BoxAnnotator()images[]titles[]forkeyinkeys:frame_with_annotationsbox_annotator.annotate(sceneds.images[key].copy(),detectionsds.annotations[key],skip_labelTrue)images.append(frame_with_annotations)titles.append(annotations)frame_with_predictionsbox_annotator.annotate(sceneds.images[key].copy(),detectionspredictions[key],skip_labelTrue)images.append(frame_with_predictions)titles.append(predictions)%matplotlib inline sv.plot_images_grid(imagesimages,titlestitles,grid_size(n,2),size(2*4,n*4)) 结果评估最后补上混淆矩阵整体看一遍分类和检测效果。!pip install onemetricimportosimportnumpyasnpfromonemetric.cv.object_detectionimportConfusionMatrix keyslist(ds.images.keys())annotation_batches,prediction_batches[],[]forkeyinkeys:annotationds.annotations[key]annotation_batchnp.column_stack((annotation.xyxy,annotation.class_id))annotation_batches.append(annotation_batch)predictionpredictions[key]prediction_batchnp.column_stack((prediction.xyxy,prediction.class_id,prediction.confidence))prediction_batches.append(prediction_batch)confusion_matrixConfusionMatrix.from_detections(true_batchesannotation_batches,detection_batchesprediction_batches,num_classeslen(ds.classes),conf_thresholdCONFIDENCE_TRESHOLD)confusion_matrix.plot(os.path.join(HOME,confusion_matrix.png),class_namesds.classes) 小结YOLO-NAS 的关键在于 SuperGradients 的训练参数和数据路径。只要数据集格式、类别列表和 checkpoint 路径对齐就能顺利完成训练和测试。这一类 notebook 建议按“先环境、再数据、再单样例、最后批量推理”的顺序复现。遇到报错时优先检查 GPU、依赖版本、数据集目录和模型权重路径。后续我会继续按源项目顺序整理同系列中的目标检测、实例分割、OCR、多目标跟踪和视觉大模型教程。 同系列教程汇总Google Gemini 3.5 Flash 零样本目标检测教程从提示词到可视化结果GLM-OCR 文档识别实战教程从验证码、公式到车牌 OCRRF-DETR ByteTrack 多目标跟踪实战教程从命令行到 Python 视频轨迹可视化SAM 3 图像分割实战教程文本、框和点提示的多种分割方式YOLO-NAS 自定义目标检测训练实战SuperGradients 训练与测试
