环境配置情况
torch=1.9.0+cu111
torchvision=0.10.0+cu111
我的是3060按照上述torch配置就行,其余就按照yolo官方的requirement自行配置。
数据集导入
./data/coco.yaml
以这个数据集作为模板去改
参数配置如下
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../MyDataset/yolo # dataset root dir
train: images/train # train images (relative to 'path') 118287 images
val: images/val # val images (relative to 'path') 5000 images
test: images/test # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
# Classes
names:
1: hp_cm
2: hp_cd
3: kp
注意,path
的参数是直接数据集根目录,其文件结构是,如果可能请按照这个目录结果来放置数据集
# parent
# ├── yolov5
# └── MyDataset
# └── yolo
# └── images (这里放jpg)
# └── train
# └── val
# └── labels
# └── train
# └── xxxx.txt
..........
# └── val
网络模型选型
在models/yolov5x.yaml
中修改参数。(ps:我这里是选择使用yolov5x这结构)
只需要改这个参数就行
nc:3 # number of classes
运行命令
python train.py --data ./data/defect.yaml --cfg yolov5x.yaml --weight ./pretrain/yolov5x.pt --batch-size 8 --workers 0
导入的路径不需要加引号
输出的结果
最后输出的结果会在runs/train
下的每个exp中
测试集验证
运行命令
python val.py --data ./data/defect.yaml --weight ./code/yolov5/runs/train/exp5/weights/best.pt --batch-size 8 --workers 0