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 训练自己的数据集