【YOLO_NAS训练自己的数据集】

yolo格式数据集准备

数据标注由labelimg完成(共34个类别)，导出后为VOC xml格式，每张图的标注信息保存在Annotations文件夹内，后统一转化成yolo格式，数据集文件夹下有images和labels文件夹，每个文件夹内分别有train,val文件夹存放png图片和txt格式的标注信息，如图所示：

配置环境

先创建虚拟环境，
Python 3.7, 3.8 or 3.9 .
1.9.0 <= torch < 1.14 .
个人安装了python3.9、torch1.10、cuda11.3、cudnn8.2
官方github: https://github.com/Deci-AI/super-gradients

安装有两种方式：
第一种：

pip install super-gradients

第二种：

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

根据自己数据集修改配置文件

1.先说训练命令（training with a recipe）

第一种（官方示例）：

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

这里的运行的文件是包里面的src/super_gradients/train_from_recipe.py文件，后面指定一个配置yaml文件用来设置训练相关参数。

我的命令则是如下：

python -m super_gradients.train_from_recipe --config-name=ronghuahanjie_yolo_nas_s.yaml

第二种：
进入super_gradients目录，执行

python src/super_gradients/examples/train_from_recipe_example/train_from_recipe.py --config-name=ronghuahanjie_yolo_nas_s.yaml

2.修改训练脚本主配置文件

位置：src/super_gradients/recipes
复制一份coco2017_yolo_nas_s.yaml文件放在同一文件夹，命名成自己的配置文件并修改，修改示例如下：

名字根据自己项目名字重新命名就行，然后依次修改对应的文件就行。
注：
我自己只用一块12G显存的GPU，且后续训练验证出现DDP引起的报错中断，程序被Kill，因此我这里把DDP分布式模式关了；最后一项实验名字，只要在训练周期保持一致就行，中途不要随意改动，设置什么名字无所谓，把coco2017_替换成自己的也行。

3.修改dataset_params

复制一份coco_detection_yolo_format_base_dataset_params.yaml 并重命名，位置在
/src/super_gradients/recipes/dataset_params/ronghuahanjie_detection_yolo_format_base_dataset_params.yaml

把文件内训练验证数据集位置和标签换成自己的就行

4.修改dataloaders

修改src/super_gradients/training/dataloaders/dataloaders.py
复制coco_detection_yolo_format_train、coco_detection_yolo_format_val函数并修改成自己的：

5.修改object_names

修改src/super_gradients/common/object_names.py,增添自己的

训练

接下来就可以训练了，训练上面说过了，这里补充下中途中断训练怎么恢复，如下

python src/super_gradients/examples/train_from_recipe_example/train_from_recipe.py --config-name=ronghuahanjie_yolo_nas_s.yaml training_hyperparams.resume=True

训练的权重等信息默认保存在super_gradients/checkpoints中，恢复训练默认从最近的ckpt_latest.pth开始；如果你不想从最近的checkpoint恢复训练，指定文件的话，通过指定resume_path实现，且当使用resume_path时不需要设置 resume=True
该相关issue反馈链接：添加链接描述

推理

import os
import torch
from super_gradients.training import models
from super_gradients.common.object_names import Models

# 加载预训练模型
net = models.get(Models.YOLO_NAS_S, pretrained_weights="coco")
# 加载自己训练的模型
net = models.get(Models.YOLO_NAS_S, num_classes=34, checkpoint_path=os.path.join(path, "ckpt_best.pth"))

net = net.to("cuda" if torch.cuda.is_available() else "cpu")
net.eval()


# 1.推理在线图片
url = "https://www.aljazeera.com/wp-content/uploads/2022/12/2022-12-03T205130Z_851430040_UP1EIC31LXSAZ_RTRMADP_3_SOCCER-WORLDCUP-ARG-AUS-REPORT.jpg?w=770&resize=770%2C436&quality=80"
prediction = net.predict(url, conf=0.25)
prediction.show()

# 2.推理本地图片
prediction = net.predict("xx/xxxxxx.png")
prediction.show()

# 3.推理本地视频
net.predict(input_video_path).save(output_video_path)

# 4.webcam实时推理
net.predict_webcam()

# 5.转化onnx格式
net.prep_model_for_conversion(input_size=[1, 3, 640, 640])
# Create dummy_input
# Convert model to onnx
torch.onnx.export(net, dummy_input,  "yolo_nas_s.onnx")

最后感谢大佬的指南引路 ：YOLO NAS 训练自己的数据集