super-gradients

构建、训练和微调生产就绪的深度学习SOTA视觉模型

3.5版本已发布！笔记本已更新！

</div> <div align="center"> <p align="center"> <a href="#getting-started">入门</a> • <a href="#implemented-model-architectures">预训练模型</a> • <a href="#community">社区</a> • <a href="#license">许可证</a> • </p> <p align="center"> <a href="https://github.com/Deci-AI/super-gradients#prerequisites"><img src="https://img.shields.io/badge/python-3.7%20%7C%203.8%20%7C%203.9-blue" /></a> <a href="https://github.com/Deci-AI/super-gradients#prerequisites"><img src="https://img.shields.io/badge/pytorch-1.9%20%7C%201.10-blue" /></a> <a href="https://pypi.org/project/super-gradients/"><img src="https://img.shields.io/pypi/v/super-gradients" /></a> <a href="https://github.com/Deci-AI/super-gradients/blob/master/documentation/source/model_zoo.md" ><img src="https://img.shields.io/badge/pre--trained%20models-34-brightgreen" /></a> <a href="https://github.com/Deci-AI/super-gradients/releases"><img src="https://img.shields.io/github/v/release/Deci-AI/super-gradients" /></a> <a href="https://join.slack.com/t/supergradients-comm52/shared_invite/zt-10vz6o1ia-b_0W5jEPEnuHXm087K~t8Q"><img src="https://img.shields.io/badge/slack-community-blueviolet" /></a> <a href="https://github.com/Deci-AI/super-gradients/blob/master/LICENSE.md"><img src="https://img.shields.io/badge/license-Apache%202.0-blue" /></a> <a href="https://docs.deci.ai/super-gradients/documentation/source/welcome.html"><img src="https://img.shields.io/badge/docs-mkdocs-brightgreen" /></a> </p> </div>

使用SuperGradients构建

支持各种计算机视觉任务

可直接部署的预训练SOTA模型

YOLO-NAS和YOLO-NAS-POSE架构已发布！新的YOLO-NAS以无与伦比的准确度-速度性能表现出色，超越了其他模型如YOLOv5、YOLOv6、YOLOv7和YOLOv8。用于姿态估计的YOLO-NAS-POSE模型也已推出，提供了最先进的准确度/性能权衡。

在这里查看详情：YOLO-NAS和YOLO-NAS-POSE。

# 加载带有预训练权重的模型
from super_gradients.training import models
from super_gradients.common.object_names import Models

model = models.get(Models.YOLO_NAS_M, pretrained_weights="coco")

所有计算机视觉模型的预训练检查点可在模型库中找到

分类

语义分割

目标检测

姿态估计

轻松训练SOTA模型

轻松加载和微调生产就绪的预训练SOTA模型，这些模型结合了最佳实践和经过验证的超参数，以实现一流的准确性。有关如何操作的更多信息，请参阅入门

即插即用的配方

python -m super_gradients.train_from_recipe architecture=regnetY800 dataset_interface.data_dir=<YOUR_Imagenet_LOCAL_PATH> ckpt_root_dir=<CHEKPOINT_DIRECTORY>

关于如何以及为什么使用配方的更多示例可以在配方中找到

生产就绪

所有SuperGradients模型都已生产就绪，因为它们与TensorRT（Nvidia）和OpenVINO（Intel）等部署工具兼容，可以轻松投入生产。只需几行代码，您就可以轻松地将模型集成到您的代码库中。

# 加载带有预训练权重的模型
from super_gradients.training import models
from super_gradients.common.object_names import Models

model = models.get(Models.YOLO_NAS_M, pretrained_weights="coco")

# 准备模型进行转换
# 输入大小格式为[Batch x Channels x Width x Height]，其中640是标准COCO数据集尺寸
model.eval()
model.prep_model_for_conversion(input_size=[1, 3, 640, 640])
    
# 创建虚拟输入

# 将模型转换为onnx
torch.onnx.export(model, dummy_input,  "yolo_nas_m.onnx")

关于如何将模型投入生产的更多信息可以在入门笔记本中找到

快速安装

pip install super-gradients

入门

只需1条命令行即可开始训练

使用SuperGradients可重现的配方训练SOTA性能模型的最简单直接方法。只需定义数据集路径和保存检查点的位置,就可以在终端中开始训练了!

只需确保按照配方中指定的数据目录设置好数据集。

python -m super_gradients.train_from_recipe --config-name=imagenet_regnetY architecture=regnetY800 dataset_interface.data_dir=<你的Imagenet本地路径> ckpt_root_dir=<检查点目录>

快速加载预训练权重以获得SOTA性能的所需模型

想在自己的机器上尝试我们的预训练模型吗?导入SuperGradients,初始化你的Trainer,并从我们的SOTA模型库加载所需的架构和预训练权重

# pretrained_weights参数将加载在提供的数据集上预训练的架构

import super_gradients

model = models.get("模型名称", pretrained_weights="预训练模型名称")

语义分割

姿态估计

目标检测

如何使用预训练模型进行预测

如何使用预训练模型进行预测

Albumentations集成

在SG中使用Albumentations

高级功能

训练后量化和量化感知训练

量化涉及以低精度表示权重和偏置，从而减少内存和计算需求，这对于在资源有限的设备上部署模型很有用。这个过程可以在训练期间进行，称为量化感知训练，或在训练后进行，称为训练后量化。完整教程可以在这里找到。

训练后量化和量化感知训练

在自定义数据集上进行YoloNAS量化感知训练

本教程提供了如何使用自定义数据集微调YoloNAS模型的全面指南。它还演示了如何利用SG的QAT（量化感知训练）支持。此外，它还提供了部署模型和进行基准测试的分步说明。

在自定义数据集上进行YoloNAS量化感知训练

知识蒸馏训练

知识蒸馏是一种训练技术，使用大型模型（教师模型）来提高小型模型（学生模型）的性能。通过我们在Google Colab上使用预训练的BEiT base教师模型和Resnet18学生模型在CIFAR10上的示例notebook，了解更多关于SuperGradients知识蒸馏训练的信息，这是一个易于使用的教程，使用免费的GPU硬件。

知识蒸馏训练

配方

要训练模型，需要配置4个主要组件。这些组件被聚合到一个"主要"配方.yaml文件中，该文件继承了上述数据集、架构、训练和检查点参数。也可以（并建议为了灵活性）用自定义设置覆盖默认设置。所有配方都可以在这里找到 </br> 配方支持在SuperGradients中实现的每个模型、指标或损失函数，但你可以通过"注册"轻松扩展到任何你需要的自定义对象。查看这个教程了解更多信息。

如何使用配方

<details markdown="1"> <summary><h3>使用分布式数据并行(DDP)</h3></summary>

为什么使用DDP？

最近的深度学习模型越来越大，以至于在单个GPU上训练可能需要数周时间。为了及时训练模型，有必要使用多个GPU进行训练。使用100多个GPU可以将模型的训练时间从一周减少到不到一小时。

它是如何工作的？

每个GPU都有自己的进程，该进程控制模型的一个副本，并在训练期间从磁盘加载自己的小批量数据并将其发送到GPU。在每个GPU上完成前向传播后，梯度在所有GPU之间进行归约，导致所有GPU在本地具有相同的梯度。这使得在反向传播后，模型权重在所有GPU上保持同步。

如何使用它？

你可以使用SuperGradients在几行代码内用DDP训练你的模型。

main.py

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_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL, num_gpus=4)

# 调用训练器
Trainer(expriment_name=...)

# 你在下面做的一切都会在4个gpu上运行

...

Trainer.train(...)

最后，你可以用一个简单的python调用来启动你的分布式训练。

python main.py

请注意，如果你使用torch<1.9.0（已弃用），你将需要使用torch.distributed.launch或torchrun来启动你的训练，在这种情况下，nproc_per_node将覆盖gpu_mode设置的值：

python -m torch.distributed.launch --nproc_per_node=4 main.py

torchrun --nproc_per_node=4 main.py

在单个节点上调用函数

在DDP训练中，我们经常希望在主节点（即rank 0）上执行代码。在SG中，用户通常通过触发"阶段回调"来执行自己的代码（参见下面的"使用阶段回调"部分）。可以使用ddp_silent_mode或multi_process_safe装饰器确保所需的代码只在rank 0上运行。例如，考虑下面这个简单的阶段回调，它在训练期间将每个批次的前3张图像上传到Tensorboard：

from super_gradients.training.utils.callbacks import PhaseCallback, PhaseContext, Phase
from super_gradients.common.environment.env_helpers import multi_process_safe

class Upload3TrainImagesCalbback(PhaseCallback):
    def __init__(
        self,
    ):
        super().__init__(phase=Phase.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(tag=tag, images=batch_imgs[: 3], global_step=context.epoch)

@multi_process_safe 装饰器确保回调只会被 rank 0 触发。另外,也可以通过 SG trainer 的布尔属性 ddp_silent_mode 来实现(phase context 可以访问该属性)。该属性仅当当前进程 rank 为零时设为 False(即使进程组已被销毁):

from super_gradients.training.utils.callbacks import PhaseCallback, PhaseContext, Phase

class Upload3TrainImagesCalbback(PhaseCallback):
    def __init__(
        self,
    ):
        super().__init__(phase=Phase.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(tag=tag, images=batch_imgs[: 3], global_step=context.epoch)

注意 ddp_silent_mode 可以通过 SgTrainer.ddp_silent_mode 访问。因此,在调用 SgTrainer.train() 后,如果脚本的某些部分只需在 rank 0 上运行,可以使用它。

值得了解

你的总批量大小将是 (GPU 数量 x 批量大小),所以你可能需要增加学习率。没有明确的规则,但经验法则似乎是根据 GPU 数量线性增加学习率

</details> <details markdown="1"> <summary><h3> 轻松更改架构参数 </h3></summary>

from super_gradients.training import models

# 实例化默认预训练 resnet18
default_resnet18 = models.get(model_name="resnet18", num_classes=100, pretrained_weights="imagenet")

# 实例化预训练 resnet18,开启 DropPath 并设置概率为 0.5
droppath_resnet18 = models.get(model_name="resnet18", arch_params={"droppath_prob": 0.5}, num_classes=100, pretrained_weights="imagenet")

# 实例化预训练 resnet18,不含分类器头。输出将来自全局池化前的最后一个阶段
backbone_resnet18 = models.get(model_name="resnet18", arch_params={"backbone_mode": True}, pretrained_weights="imagenet")

</details> <details markdown="1"> <summary><h3> 使用阶段回调 </h3></summary>

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 训练和验证数据加载器以及优化器

# 定义回调中要调用的内容
rop_lr_scheduler = ReduceLROnPlateau(optimizer, mode="max", patience=10, verbose=True)

# 定义阶段回调,它们将按 Phase 中定义的方式触发
phase_callbacks = [LRSchedulerCallback(scheduler=rop_lr_scheduler,
                                       phase=Phase.VALIDATION_EPOCH_END,
                                       metric_name="Accuracy")]

# 创建一个训练器对象,查看声明以获取更多参数
trainer = Trainer("experiment_name")

# 将 phase_callbacks 定义为训练参数的一部分
train_params = {"phase_callbacks": phase_callbacks}

</details> <details markdown="1"> <summary><h3> 与 DagsHub 集成 </h3></summary>

from super_gradients import Trainer

trainer = Trainer("experiment_name")
model = ...

training_params = { ...  # 你的训练参数
                   "sg_logger": "dagshub_sg_logger",  # DagsHub 日志记录器,详见 super_gradients.common.sg_loggers.dagshub_sg_logger.DagsHubSGLogger 类
                   "sg_logger_params":  # 将传递给日志记录器 super_gradients.common.sg_loggers.dagshub_sg_logger.DagsHubSGLogger 的 __init__ 的参数
                     {
                       "dagshub_repository": "<REPO_OWNER>/<REPO_NAME>", # 可选: 你的 DagsHub 项目名称,由所有者名称、'/'和仓库名称组成。如果留空,你将在运行时被提示手动填写。
                       "log_mlflow_only": False, # 可选: 改为 true 以绕过 DVC 日志记录,仅将所有工件记录到 MLflow  
                       "save_checkpoints_remote": True,
                       "save_tensorboard_remote": True,
                       "save_logs_remote": True,
                     }
                   }

</details> <details> <summary><h3> 与 Weights and Biases 集成 </h3></summary>

from super_gradients import Trainer

# 创建一个训练器对象,查看声明以获取更多参数
trainer = Trainer("experiment_name")

train_params = { ... # 训练参数
                "sg_logger": "wandb_sg_logger", # Weights&Biases 日志记录器,详见 WandBSGLogger 类
                "sg_logger_params": # 将传递给日志记录器 __init__ 的参数 
                  {
                    "project_name": "project_name", # W&B 项目名称
                    "save_checkpoints_remote": True
                    "save_tensorboard_remote": True
                    "save_logs_remote": True
                  } 
               }

</details> <details markdown="1"> <summary><h3> 与 ClearML 集成 </h3></summary>

from super_gradients import Trainer

# 创建一个训练器对象,查看声明以获取更多参数
trainer = Trainer("experiment_name")

train_params = { ... # 训练参数
                "sg_logger": "clearml_sg_logger", # ClearML 日志记录器,详见 ClearMLSGLogger 类
                "sg_logger_params": # 将传递给日志记录器 __init__ 的参数 
                  {
                    "project_name": "project_name", # ClearML 项目名称
                    "save_checkpoints_remote": True,
                    "save_tensorboard_remote": True,
                    "save_logs_remote": True,
                  } 
               }

</details> <details markdown="1"> <summary><h3> 与 Voxel51 集成 </h3></summary>

你可以使用 apply_model() 方法直接将 SuperGradients YOLO-NAS 模型应用于你的 FiftyOne 数据集:

import fiftyone as fo
import fiftyone.zoo as foz

from super_gradients.training import models

dataset = foz.load_zoo_dataset("quickstart", max_samples=25)
dataset.select_fields().keep_fields()

model = models.get("yolo_nas_m", pretrained_weights="coco")

dataset.apply_model(model, label_field="yolo_nas", confidence_thresh=0.7)

session = fo.launch_app(dataset)

可以直接从 FiftyOne Model Zoo 访问 SuperGradients YOLO-NAS 模型:

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_field="yolo_nas")

session = fo.launch_app(dataset)

</details>

安装方法

前提条件

已安装 Python 3.7、3.8 或 3.9。
1.9.0 <= torch < 1.14
- https://pytorch.org/get-started/locally/
requirements.txt 中指定的 Python 包;

</details> <details markdown="1"> <summary>在 nvidia GPU 上训练</summary>

Nvidia CUDA Toolkit >= 11.2
CuDNN >= 8.1.x
支持 CUDA >= 11.2 的 Nvidia 驱动程序 (≥460.x)

</details>

快速安装

<details markdown="1"> <summary>使用 PyPi 安装稳定版本</summary>

参见 PyPi

pip install super-gradients

就这么简单!

</details> <details markdown="1"> <summary>使用 GitHub 安装</summary>

pip install git+https://github.com/Deci-AI/super-gradients.git@stable

</details>

已实现的模型架构

所有计算机视觉模型 - 预训练检查点可在模型库中找到

图像分类

语义分割

目标检测

姿态估计

DEKR-W32-NO-DC

已实现的数据集

Deci提供了各种数据集的实现。如果您需要下载任何数据集，可以查看说明。

图像分类

语义分割

目标检测

姿态估计

COCO

文档

查看SuperGradients文档以获取完整文档、用户指南和示例。

贡献

要了解如何为SuperGradients做出贡献，请参阅我们的贡献页面。

我们出色的贡献者：

<br/>由contrib.rocks制作。

引用

如果您在研究中使用SuperGradients库或基准测试，请引用SuperGradients深度学习训练库。

社区

如果您想成为SuperGradients不断壮大的社区的一员，了解所有激动人心的新闻和更新，需要帮助，请求高级功能，或想要报告错误或问题，我们非常欢迎您的加入！

Discord是讨论SuperGradients并获得支持的地方。点击此处加入我们的Discord社区
要报告错误，请在GitHub上提交问题。
加入SG通讯，了解新功能和模型、重要公告以及即将举行的活动。
如需与我们进行简短会面，请使用此链接并选择您的首选时间。

许可证

本项目根据Apache 2.0许可证发布。

引用

BibTeX

@misc{supergradients,
  doi = {10.5281/ZENODO.7789328},
  url = {https://zenodo.org/record/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},
}