YOLO NAS 训练自己的数据集
环境准备
yolo nas源码:https://github.com/Deci-AI/super-gradients.git
参考readme文件:环境配置一行指令即可
pip install super-gradients
数据集准备
yolo nas目标检测支持coco,voc,yolo三种格式的数据集,在这里使用yolo格式的数据集.
如 super-gradients/src/super_gradients/training/datasets/detection_datasets
下文件有说明:
文件修改
本人是葡萄数据集,所以文件中以grape命名.
传入自己的数据集需要修改以下文件:
1.recipes主目录下
目录:super-gradients/src/super_gradients/recipes/grape_yolo_nas_s.yaml
复制coco_yolo_nas_s.yaml,修改为自己的配置文件,其它文件改名同理:
defaults:
- training_hyperparams: coco2017_yolo_nas_train_params
- # 文件目录src/super_gradients/recipes/training_hyperparams/coco2017_yolo_nas_train_params.yaml
- # 需要修改epochs在此文件中修改
- dataset_params: grape_detection_yolo_format_base_dataset_params
- # 文件目录src/super_gradients/recipes/dataset_params/grape_detection_yolo_format_base_dataset_params.yaml
- # 此处修改见说明-dataset_params修改
- arch_params: yolo_nas_s_arch_params
- checkpoint_params: default_checkpoint_params
- _self_
- variable_setup
train_dataloader: grape_detection_yolo_format_train
val_dataloader: grape_detection_yolo_format_val
load_checkpoint: False
resume: False
dataset_params:
train_dataloader_params:
batch_size: 4
arch_params:
num_classes: 1
training_hyperparams:
resume: ${
resume}
mixed_precision: True
architecture: yolo_nas_s
multi_gpu: DDP
num_gpus: 1
experiment_suffix: ""
experiment_name: coco2017_${
architecture}${
experiment_suffix}
-dataset_params修改:
目录/src/super_gradients/recipes/dataset_params/grape_detection_yolo_format_base_dataset_params.yaml
1.训练数据集目录修改
2.验证数据集目录修改
2.dataloaders.py文件修改
文件目录:src/super_gradients/training/dataloaders/dataloaders.py
复制coco_detection_yolo_format_train函数并修改:
复制coco_detection_yolo_format_val函数并修改:
注意:此处config_name命名对应的文件是步骤一中dataset_params修改的文件除去后缀名
3.object_names.py文件修改
文件目录:src/super_gradients/common/object_names.py
开始训练
python src/super_gradients/examples/train_from_recipe_example/train_from_recipe.py --config-name=grape_yolo_nas_s.yaml