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

在这里插入图片描述

猜你喜欢

转载自blog.csdn.net/balabala_333/article/details/132183996