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