Project Icon

super-gradients

开源工具库简化SOTA计算机视觉模型的训练与部署

Super-Gradients是一个专注于计算机视觉的开源深度学习库。它提供预训练SOTA模型和易用训练工具,支持分类、分割、检测等任务。该项目集成多种训练技巧,兼容主流部署框架,可快速将模型应用于生产。Super-Gradients适用于学术研究和工业应用,是一个高效的计算机视觉开发工具。



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

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



使用SuperGradients构建


支持各种计算机视觉任务

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

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

在这里查看详情:YOLO-NASYOLO-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

最新动态


3.4.0版本(2023年11月6日)

  • YoloNAS-Pose模型发布 - 姿态估计的新前沿
  • 添加了将配方导出为单个YAML文件或独立train.py文件的选项
  • 其他错误修复和小改进。完整发布说明可在此处查看

3.1.3版本(2023年7月19日)


5月30日

3.1.1版本(5月3日)

查看SG完整发布说明

目录


入门


只需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集成

高级功能


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

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

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

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

知识蒸馏训练

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

配方

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

使用分布式数据并行(DDP)

为什么使用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.launchtorchrun来启动你的训练,在这种情况下,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 数量线性增加学习率

轻松更改架构参数

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")

使用阶段回调

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}

与 DagsHub 集成

Open In Colab

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,
                     }
                   }

与 Weights and Biases 集成

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
                  } 
               }

与 ClearML 集成

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,
                  } 
               }

与 Voxel51 集成

你可以使用 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)

安装方法


前提条件

一般要求
在 nvidia GPU 上训练

快速安装

使用 PyPi 安装稳定版本

参见 PyPi

pip install super-gradients

就这么简单!

使用 GitHub 安装
pip install git+https://github.com/Deci-AI/super-gradients.git@stable

已实现的模型架构


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

图像分类

语义分割

目标检测

姿态估计


已实现的数据集


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

图像分类

语义分割

目标检测

姿态估计


文档

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

贡献

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

我们出色的贡献者:


contrib.rocks制作。

引用

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

社区

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

许可证

本项目根据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},
}

最新DOI

DOI


在此处申请免费试用

项目侧边栏1项目侧边栏2
推荐项目
Project Cover

豆包MarsCode

豆包 MarsCode 是一款革命性的编程助手,通过AI技术提供代码补全、单测生成、代码解释和智能问答等功能,支持100+编程语言,与主流编辑器无缝集成,显著提升开发效率和代码质量。

Project Cover

问小白

问小白是一个基于 DeepSeek R1 模型的智能对话平台,专为用户提供高效、贴心的对话体验。实时在线,支持深度思考和联网搜索。免费不限次数,帮用户写作、创作、分析和规划,各种任务随时完成!

Project Cover

白日梦AI

白日梦AI提供专注于AI视频生成的多样化功能,包括文生视频、动态画面和形象生成等,帮助用户快速上手,创造专业级内容。

Project Cover

有言AI

有言平台提供一站式AIGC视频创作解决方案,通过智能技术简化视频制作流程。无论是企业宣传还是个人分享,有言都能帮助用户快速、轻松地制作出专业级别的视频内容。

Project Cover

讯飞绘镜

讯飞绘镜是一个支持从创意到完整视频创作的智能平台,用户可以快速生成视频素材并创作独特的音乐视频和故事。平台提供多样化的主题和精选作品,帮助用户探索创意灵感。

Project Cover

讯飞文书

讯飞文书依托讯飞星火大模型,为文书写作者提供从素材筹备到稿件撰写及审稿的全程支持。通过录音智记和以稿写稿等功能,满足事务性工作的高频需求,帮助撰稿人节省精力,提高效率,优化工作与生活。

Project Cover

阿里绘蛙

绘蛙是阿里巴巴集团推出的革命性AI电商营销平台。利用尖端人工智能技术,为商家提供一键生成商品图和营销文案的服务,显著提升内容创作效率和营销效果。适用于淘宝、天猫等电商平台,让商品第一时间被种草。

Project Cover

Trae

Trae是一种自适应的集成开发环境(IDE),通过自动化和多元协作改变开发流程。利用Trae,团队能够更快速、精确地编写和部署代码,从而提高编程效率和项目交付速度。Trae具备上下文感知和代码自动完成功能,是提升开发效率的理想工具。

Project Cover

AIWritePaper论文写作

AIWritePaper论文写作是一站式AI论文写作辅助工具,简化了选题、文献检索至论文撰写的整个过程。通过简单设定,平台可快速生成高质量论文大纲和全文,配合图表、参考文献等一应俱全,同时提供开题报告和答辩PPT等增值服务,保障数据安全,有效提升写作效率和论文质量。

投诉举报邮箱: service@vectorlightyear.com
@2024 懂AI·鲁ICP备2024100362号-6·鲁公网安备37021002001498号