)
SuperGradients 实战指南一站式训练、微调与部署 SOTA 视觉模型含 YOLO-NAS 生态【免费下载链接】super-gradientsEasily train or fine-tune SOTA computer vision models with one open source training library. The home of Yolo-NAS.项目地址: https://gitcode.com/GitHub_Trending/su/super-gradientsSuperGradientsSG是 Deci 开源的一站式计算机视觉训练库它把图像分类、语义分割、目标检测与姿态估计四大任务的模型训练、预训练权重加载、微调、量化与生产部署统一到同一套 API 与配置体系之下也是 YOLO-NAS 与 YOLO-NAS-POSE 系列模型的官方开源载体。读完本文你将掌握如何用一条命令行启动 SOTA 模型的训练、用几行 Python 加载预训练权重完成推理与模型导出并理解 Recipes 配置、分布式训练DDP、Phase Callbacks 与实验监控集成等进阶能力的底层实现。SuperGradients 是什么一个库覆盖四大 CV 任务SuperGradients 的核心定位是构建、训练并微调生产就绪的 SOTA 视觉模型仓库 README.md 首页标语。它并非只封装某一种模型而是围绕Trainer统一了完整的训练管线数据加载、数据增强、损失函数、指标、学习率调度、EMA、混合精度、量化训练QAT、知识蒸馏KD等先进训练技巧全部开箱即用同时保证训练出的模型可直接部署到 TensorRT、OpenVINO 等推理引擎。围绕四大视觉任务SuperGradients 分别提供了开箱即用的实现与预训练权重图像分类ResNet、RegNet、EfficientNet、MobileNet v2/v3、DenseNet、RepVGG、ViT 等语义分割DDRNet、ShelfNet、STDC、PP-LiteSeg、RegSeg、SegFormer 等目标检测YOLO-NASS/M/L、YOLOX、SSD、PP-YOLOE 等姿态估计YOLO-NAS-POSEN/S/M/L、DEKR 等。其中 YOLO-NAS 是 Deci 基于其 AutoNAC™ 神经架构搜索技术生成的下一代目标检测模型据仓库 YOLONAS.md 文档所述其在 COCO 2017 Val 上以更优的精度/速度权衡对标 YOLOv5/v6/v7/v8 等模型YOLO-NAS-POSE 则将该系列扩展到姿态估计任务详见 YOLONAS-POSE.md。快速安装与环境要求前置条件根据 README.md 的 Prerequisites 说明仓库要求Python 3.7 / 3.8 / 3.91.9.0 torch 1.14PyTorch 请按官方指引安装其余依赖见仓库根目录 requirements.txt。若要使用 NVIDIA GPU 训练还需要 CUDA Toolkit ≥ 11.2、CuDNN ≥ 8.1.x以及支持 CUDA ≥ 11.2 的驱动≥ 460.x。注当前仓库版本为 3.7.1见 src/super_gradients/init.pyREADME 顶部横幅提到的 Version 3.5 is out 为历史发布公告仓库内的 notebook 均已同步更新。安装方式从 PyPI 安装稳定版pip install super-gradients也可以直接从 Git 安装仅用于说明安装方式git clone https://gitcode.com/GitHub_Trending/su/super-gradients cd super-gradients pip install .安装完成后即可在任意 Python 脚本中import super_gradients。一条命令行启动 SOTA 模型训练Recipes 机制SuperGradients 最具特色的能力是用一条命令行训练 SOTA 模型其底层是 Hydra 驱动的 Recipes 配置体系。一行命令训练 RegNetYImageNet 分类在仓库根目录执行以 imagenet_regnetY.yaml 为例python -m super_gradients.train_from_recipe --config-nameimagenet_regnetY architectureregnetY800 dataset_interface.data_dirYOUR_Imagenet_LOCAL_PATH ckpt_root_dirCHEKPOINT_DIRECTORY你只需要提供--config-name要使用的 recipe 名、architecture具体网络结构如regnetY200/400/600/800、dataset_interface.data_dir数据集路径与ckpt_root_dir权重保存目录。数据集需按 recipe 中指定的目录结构预先放置各数据集准备说明见 src/super_gradients/training/datasets 下的说明文档。入口实现train_from_recipe 是怎么工作的python -m super_gradients.train_from_recipe对应的入口在 src/super_gradients/train_from_recipe.pyhydra.main(config_pathrecipes, version_base1.2) def _main(cfg: DictConfig) - None: Trainer.train_from_config(cfg)即Hydra 从recipes/目录加载配置最终交给Trainer.train_from_config执行完整训练流程。Recipe 的组成四个核心配置组件一个主 recipe.yaml文件聚合了训练所需的四大组件README.md Recipes 一节并通过defaults继承子配置training_hyperparams训练超参数如 batch size、学习率、loss、metric、回调dataset_params数据集与增强配置arch_params网络结构参数checkpoint_params断点与权重加载配置。以 imagenet_regnetY.yaml 为例其defaults段如下defaults: - training_hyperparams: imagenet_regnetY_train_params - dataset_params: imagenet_regnetY_dataset_params - arch_params: regnetY_arch_params - checkpoint_params: default_checkpoint_params - _self_ - variable_setup主 recipe 还允许就地覆写默认值例如指定num_classes、dropout_prob、droppath_prob以及通过architecture: regnetY800指定具体结构。所有内置 recipe 均位于 src/super_gradients/recipes覆盖分类、检测、分割、姿态估计等任务官方还提供 Roboflow-100 数据集上的 YOLO-NAS 微调 recipe如 roboflow_yolo_nas_m.yaml。想要更深入理解 recipe 的用法可参考 notebook what_are_recipes_and_how_to_use.ipynb以及文档 Recipes_Custom.md、Recipes_Training.md。几行代码加载预训练 SOTA 模型使用 models.get 加载权重from super_gradients.training import models from super_gradients.common.object_names import Models model models.get(Models.YOLO_NAS_M, pretrained_weightscoco)models.get是加载任意内置架构的核心入口其完整签名见 src/super_gradients/training/models/model_factory.py支持model_name架构名可用 src/super_gradients/common/object_names.py 中Models枚举常量如Models.YOLO_NAS_M、Models.RESNET18、Models.DDRNET_23也可直接传字符串arch_params架构超参如{droppath_prob: 0.5}num_classes类别数若不传则从预训练权重对应数据集自动推导pretrained_weights预训练权重对应的数据集名如coco、imagenetcheckpoint_path加载外部 checkpoint支持本地路径、相对路径或 URL并自动完成权重加载load_backbone仅将 checkpoint 加载到 backbone用于迁移学习checkpoint_num_classes当 checkpoint 的类别数与目标num_classes不一致时先按 checkpoint 类别数实例化并加载权重再调用replace_head(new_num_classesnum_classes)替换分类头num_input_channels修改输入通道数如将 3 通道模型改为 1 通道底层通过replace_input_channels实现。源码中instantiate_modelmodel_factory.py负责从注册表ARCHITECTURES解析架构类 → 组装arch_params→ 若指定了预训练权重则调用load_pretrained_weights加载 → 为需要预处理的模型HasPredict接口注入默认预处理参数。若模型名既不在本地注册表、也未在 Deci 平台找到会抛出UnknownTypeException并列出全部可选架构名。一行命令式推理predict 接口对于 YOLO-NAS 等带predict能力的模型可以直接对图片 URL、本地路径、视频甚至摄像头流做端到端预测详见 YOLONAS.md Quickstartimport super_gradients yolo_nas super_gradients.training.models.get(yolo_nas_l, pretrained_weightscoco).cuda() model_predictions yolo_nas.predict(https://.../sample.jpg).show() prediction model_predictions[0].prediction # 单张图片取第一个预测结果 bboxes prediction.bboxes_xyxy # [[Xmin, Ymin, Xmax, Ymax], ...] class_names prediction.class_names # 类别名列表 class_name_indexes prediction.labels.astype(int) # 每个框对应的类别索引 confidences prediction.confidence.astype(float) # 每个框的置信度预测能力的完整使用方式可参考 notebook how_to_run_model_predict.ipynb。四大任务的动手实践Notebook 索引README 为每个任务都提供了可一键运行的 Colab notebook仓库内对应文件位于 notebooks分类transfer_learning_classification.ipynb分类迁移学习、PTQ_and_QAT_for_classification.ipynb分类模型的 PTQ 与 QAT语义分割quickstart_segmentation.ipynb分割快速上手、transfer_learning_semantic_segmentation.ipynb分割迁移学习、segmentation_connect_custom_dataset.ipynb接入自定义数据集、Segmentation_Model_Export.ipynb分割模型导出 ONNX姿态估计YoloNAS_Pose_Fine_Tuning_Animals_Pose_Dataset.ipynbYOLO-NAS-Pose 微调、DEKR_PoseEstimationFineTuning.ipynbDEKR 微调目标检测YoloNAS_Inference_using_TensorRT.ipynbYOLO-NAS 的 TensorRT 推理、detection_transfer_learning.ipynb检测迁移学习、detection_how_to_connect_custom_dataset.ipynb接入自定义检测数据集、yolo_nas_custom_dataset_fine_tuning_with_qat.ipynb自定义数据集上的 QAT 微调数据增强albumentations_tutorial.ipynbSG 与 Albumentations 的集成使用。生产就绪一行代码导出 ONNXSuperGradients 的模型在设计上兼顾训练与部署通过 module_interfaces 定义统一的导出接口支持导出到 ONNX并可进一步转 TensorRT / OpenVINO 等引擎。README 给出的导出示例from super_gradients.training import models from super_gradients.common.object_names import Models model models.get(Models.YOLO_NAS_M, pretrained_weightscoco) # 准备转换输入尺寸格式为 [Batch x Channels x Width x Height]640 是 COCO 标准尺寸 model.eval() model.prep_model_for_conversion(input_size[1, 3, 640, 640]) # 创建 dummy_input 后导出 ONNX torch.onnx.export(model, dummy_input, yolo_nas_m.onnx)完整导出流程含 NMS 融合、TensorRT 部署、QAT 导出可查阅 models_export.md、models_export_pose.md、ptq_qat.md 等文档以及 notebook Segmentation_Model_Export.ipynb。高级特性量化、蒸馏与训练技巧Post Training QuantizationPTQ与 Quantization Aware TrainingQAT量化将权重与偏置用更低精度表示从而降低内存与算力需求便于部署到资源受限设备。SuperGradients 同时支持训练后量化PTQ与训练中量化QAT。YOLO-NAS 的架构本身采用量化友好设计据 YOLONAS.md 所述其 INT8 量化后的精度损失较小。相关教程见 PTQ_and_QAT_for_classification.ipynb、yolo_nas_custom_dataset_fine_tuning_with_qat.ipynb文档见 ptq_qat.md 与 qat_ptq_yolo_nas.md。Knowledge Distillation知识蒸馏蒸馏用大模型teacher提升小模型student的性能。SuperGradients 提供了内置的 KD 训练管线KDTrainer官方示例用预训练的 BEiT-base 作为 teacher、ResNet18 作为 student 在 CIFAR10 上蒸馏见 notebook how_to_use_knowledge_distillation_for_classification.ipynb文档见 KD.md。灵活修改架构参数models.get通过arch_params即可自由调整网络结构细节无需改动源码README.md Easily change architectures parameters 一节from super_gradients.training import models # 默认的预训练 resnet18100 类 default_resnet18 models.get(model_nameresnet18, num_classes100, pretrained_weightsimagenet) # 开启 DropPath概率 0.5 droppath_resnet18 models.get( model_nameresnet18, arch_params{droppath_prob: 0.5}, num_classes100, pretrained_weightsimagenet, ) # 不使用分类头输出取全局池化前的最后一层特征 backbone_resnet18 models.get( model_nameresnet18, arch_params{backbone_mode: True}, pretrained_weightsimagenet, )从源码看arch_params会被封装为HpmStruct后传入架构类构造函数model_factory.py并在实例化前通过arch_params.override(num_classesnum_classes)注入类别数。分布式训练用 DDP 在 4 张 GPU 上训练当模型大到单卡训练需要数周时多卡分布式训练成为刚需。SuperGradients 对 PyTorch DDP 做了封装几行代码即可启用。为什么用 DDPREADME 指出每个 GPU 拥有独立的进程各自持有模型副本并加载自己的 mini-batch前向完成后跨所有 GPU 做梯度归约保证各卡梯度一致从而在反向传播后各卡模型权重保持同步。如何在代码中启用from super_gradients import init_trainer, Trainer from super_gradients.common import MultiGPUMode from super_gradients.training.utils.distributed_training_utils import setup_device # 初始化环境 init_trainer() # 在 4 张 GPU 上启动 DDP setup_device(multi_gpuMultiGPUMode.DISTRIBUTED_DATA_PARALLEL, num_gpus4) # 创建 trainer其后的所有操作都会在 4 张 GPU 上执行 trainer Trainer(experiment_name...) trainer.train(...)之后直接运行python main.py注意若使用旧版torch 1.9.0已弃用需要改用torch.distributed.launch或torchrun启动此时--nproc_per_node会覆盖代码中设置的 GPU 数量python -m torch.distributed.launch --nproc_per_node4 main.py torchrun --nproc_per_node4 main.py从源码看新版Trainer构造函数已不再接受device/multi_gpu参数二者必须通过setup_device设置sg_trainer.py同时要求先调用init_trainer()完成 DDP 环境初始化否则会抛出DDPNotSetupException。仅在主进程rank 0执行代码DDP 训练中常需要只在 master rankrank 0上执行某些逻辑如上传 TensorBoard 图片。有两种方式方式一使用multi_process_safe装饰器推荐from super_gradients.training.utils.callbacks import PhaseCallback, PhaseContext, Phase from super_gradients.common.environment.env_helpers import multi_process_safe class Upload3TrainImagesCallback(PhaseCallback): def __init__(self): super().__init__(phasePhase.TRAIN_BATCH_END) multi_process_safe def __call__(self, context: PhaseContext): batch_imgs context.inputs.cpu().detach().numpy() tag batch_ str(context.batch_idx) _images context.sg_logger.add_images(tagtag, imagesbatch_imgs[:3], global_stepcontext.epoch)方式二通过context.ddp_silent_mode判断该标志仅在 rank 0 进程为 Falsefrom super_gradients.training.utils.callbacks import PhaseCallback, PhaseContext, Phase class Upload3TrainImagesCallback(PhaseCallback): def __init__(self): super().__init__(phasePhase.TRAIN_BATCH_END) def __call__(self, context: PhaseContext): if not context.ddp_silent_mode: batch_imgs context.inputs.cpu().detach().numpy() tag batch_ str(context.batch_idx) _images context.sg_logger.add_images(tagtag, imagesbatch_imgs[:3], global_stepcontext.epoch)ddp_silent_mode也可以通过Trainer.ddp_silent_mode在训练脚本中访问。DDP 训练小贴士实际 batch size GPU 数量 × 单卡 batch size因此多卡训练时通常需要相应提高学习率。README 给出的经验法则是随 GPU 数量线性增大学习率。Phase Callbacks在训练周期的任意时刻注入自定义逻辑SuperGradients 的PhaseCallback机制允许你在训练流程的特定节点执行自定义代码。可用的事件点定义在 src/super_gradients/training/utils/callbacks/base_callbacks.py 的Phase枚举中包括PRE_TRAINING训练开始前TRAIN_EPOCH_START/TRAIN_BATCH_END/TRAIN_BATCH_STEP/TRAIN_EPOCH_ENDVALIDATION_BATCH_END/VALIDATION_EPOCH_END/VALIDATION_END_BEST_EPOCHTEST_BATCH_END/TEST_ENDPOST_TRAINING训练结束后README 给出的示例用ReduceLROnPlateau在验证 epoch 结束时依据指标调整学习率from super_gradients import Trainer from torch.optim.lr_scheduler import ReduceLROnPlateau from super_gradients.training.utils.callbacks import Phase, LRSchedulerCallback from super_gradients.training.metrics.classification_metrics import Accuracy # 定义 PyTorch 的 train/val loader 与 optimizer此处省略 rop_lr_scheduler ReduceLROnPlateau(optimizer, modemax, patience10, verboseTrue) # 定义 phase callbacks在指定 Phase 触发 phase_callbacks [LRSchedulerCallback(schedulerrop_lr_scheduler, phasePhase.VALIDATION_EPOCH_END, metric_nameAccuracy)] trainer Trainer(experiment_name) # 将 callbacks 放入训练参数 train_params {phase_callbacks: phase_callbacks}更完整的回调清单见 src/super_gradients/training/utils/callbacks/callbacks.py内置了 EarlyStop、EMA、指标更新、可视化等回调也可参考 PhaseCallbacks.md。实验监控一键接入 DagsHub / WB / ClearMLSuperGradients 通过sg_logger抽象层对接主流实验管理平台实现类位于 src/super_gradients/common/sg_loggers。DagsHub 集成from super_gradients import Trainer trainer Trainer(experiment_name) model ... training_params { ... # 你的训练参数 sg_logger: dagshub_sg_logger, # 见 super_gradients.common.sg_loggers.dagshub_sg_logger.DagsHubSGLogger sg_logger_params: { dagshub_repository: REPO_OWNER/REPO_NAME, # DagsHub 项目名可留空则在运行时提示填写 log_mlflow_only: False, # 设为 True 时仅记录到 MLflow不记录到 DVC save_checkpoints_remote: True, save_tensorboard_remote: True, save_logs_remote: True, }, }Weights Biases 集成from super_gradients import Trainer trainer Trainer(experiment_name) train_params { ... # 训练参数 sg_logger: wandb_sg_logger, # 见 WandBSGLogger sg_logger_params: { project_name: project_name, # WB 项目名 save_checkpoints_remote: True, save_tensorboard_remote: True, save_logs_remote: True, }, }ClearML 集成from super_gradients import Trainer trainer Trainer(experiment_name) train_params { ... # 训练参数 sg_logger: clearml_sg_logger, # 见 ClearMLSGLogger sg_logger_params: { project_name: project_name, # ClearML 项目名 save_checkpoints_remote: True, save_tensorboard_remote: True, save_logs_remote: True, }, }Voxel51FiftyOne集成SuperGradients 的 YOLO-NAS 模型可以直接作用于 FiftyOne 数据集既能通过apply_model()直接推理也能从 FiftyOne Model Zoo 加载import fiftyone as fo import fiftyone.zoo as foz from super_gradients.training import models dataset foz.load_zoo_dataset(quickstart, max_samples25) dataset.select_fields().keep_fields() model models.get(yolo_nas_m, pretrained_weightscoco) dataset.apply_model(model, label_fieldyolo_nas, confidence_thresh0.7) session fo.launch_app(dataset)import fiftyone as fo import fiftyone.zoo as foz model foz.load_zoo_model(yolo-nas-torch) dataset foz.load_zoo_dataset(quickstart) dataset.apply_model(model, label_fieldyolo_nas) session fo.launch_app(dataset)已实现的模型架构与数据集模型架构所有内置模型的实现分布在 src/super_gradients/training/models 下的classification_models、detection_models、segmentation_models、pose_estimation_models子目录中完整清单见 src/super_gradients/training/models/Implemented Model Architectures.md可用名称的常量定义见 src/super_gradients/common/object_names.py 的Models类图像分类DenseNet、DPN、EfficientNet、LeNet、MobileNet、MobileNet v2、MobileNet v3、PNASNet、Pre-activation ResNet、RegNet、RepVGG、ResNet、ResNeXt、SENet、ShuffleNet、ShuffleNet v2、VGG语义分割PP-LiteSeg、DDRNet、LadderNet、RegSeg、ShelfNet、STDC、SegFormer目标检测CSP-DarkNet、DarkNet-53、SSD、YOLOX、PP-YOLOE、YOLO-NAS姿态估计DEKR-W32-NO-DC、YOLO-NAS-POSE。内置数据集数据集实现位于 src/super_gradients/training/datasets包括分类CIFAR-10、ImageNet分割Cityscapes、COCO、PascalVOC 2012 / PascalAUG 2012、SuperviselyPersons、Mapillary Vistas检测COCO、PascalVOC 2007 2012另支持 YOLO 格式 / Darknet 格式检测数据集对应 dataset_params 见 src/super_gradients/recipes/dataset_params姿态估计COCO。各数据集的下载与目录结构说明可参考 README.md 与数据集目录下的说明文档Detect 任务还支持 Roboflow-100 数据集集成。版本演进与生态README 的 Whats New 记录了仓库的关键能力演进3.4.0发布 YOLO-NAS-POSE 模型支持将 recipe 导出为单一 YAML 或独立train.py3.1.3加入姿态估计任务支持含预训练的 DEKR 模型、TensorRT 兼容、Python 3.10 支持、torch.compile支持3.1.1发布 YOLO-NAS、全新的predict函数支持图片、视频、URL、路径、流、Roboflow-100 集成、文档中心、DagsHub 集成、SegFormer、PTQ/QAT 支持。更多演进细节见 version.txt 与仓库 release 记录。引用与许可如果在研究中使用 SuperGradients 库或其中的基准请按仓库 README.md 提供的 BibTeX 引用misc{supergradients, doi {10.5281/ZENODO.7789328}, author {Aharon, Shay and {Louis-Dupont} and {Ofri Masad} and Yurkova, Kate and {Lotem Fridman} and {Lkdci} and Khvedchenya, Eugene and Rubin, Ran and Bagrov, Natan and Tymchenko, Borys and Keren, Tomer and Zhilko, Alexander and {Eran-Deci}}, title {Super-Gradients}, publisher {GitHub}, journal {GitHub repository}, year {2021}, }项目本体采用 Apache 2.0 许可见 LICENSE.mdYOLO-NAS 权重许可见 LICENSE.YOLONAS.mdYOLO-NAS-POSE 权重许可见 LICENSE.YOLONAS-POSE.md贡献指南见 CONTRIBUTING.md。下一步学习路径从入门到进阶建议按如下顺序阅读仓库文档完整文档中心入口documentation/source/welcome.md模型清单与权重documentation/source/model_zoo.md、documentation/source/models.md训练与配置Recipes_Training.md、Recipes_Custom.md、configuration_files.md、LRScheduling.md、EMA.md预测与部署ModelPredictions.md、PredictionSetup.md、models_export.md。至此你已经掌握了 SuperGradients 从安装、训练、微调、推理到部署的完整链路一条命令跑通 Recipe 训练、models.get加载 SOTA 预训练权重、predict做端到端推理、setup_device开启 DDP、PhaseCallback注入自定义逻辑以及通过sg_logger对接主流实验管理平台。这套以 YOLO-NAS 生态为代表的一站式方案可以直接复用到你自己的数据集与业务场景中。【免费下载链接】super-gradientsEasily train or fine-tune SOTA computer vision models with one open source training library. The home of Yolo-NAS.项目地址: https://gitcode.com/GitHub_Trending/su/super-gradients创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考