YOLOv5模型训练环境配置

记录YOLOv5使用过程中碰到的一些问题。

下载yolov5

git clone https://github.com/ultralytics/yolov5

打开cmd命令框，cd到yolov5目录下，安装依赖环境。

pip install -r requirements.txt

训练测试

python detect.py --source data\images\zidane.jpg --weights yolov5s.pt --conf 0.25

显示如下结果:

 torch-1.10.2+cpu CPU,为什么没有使用GPU？原来是pytorch的版本不对。我机器上torchvision的版本为0.11.3。

查看cuda版本

 nvcc -V

可看到CUDA的版本号为10.0。

打开requirements.txt文件，可以看到文件中要求的torch>=1.7.0 和 torchvision>=0.8.1版本。    

torch、 torchvision和python之间的版本关系。https://github.com/pytorch/vision#installation

torch	torchvision	python
main / nightly	main / nightly	>=3.7, <=3.10
1.12.0	0.13.0	>=3.7, <=3.10
1.11.0	0.12.0	>=3.7, <=3.10
1.10.2	0.11.3	>=3.6, <=3.9
1.10.1	0.11.2	>=3.6, <=3.9
1.10.0	0.11.1	>=3.6, <=3.9
1.9.1	0.10.1	>=3.6, <=3.9
1.9.0	0.10.0	>=3.6, <=3.9
1.8.2	0.9.2	>=3.6, <=3.9
1.8.1	0.9.1	>=3.6, <=3.9
1.8.0	0.9.0	>=3.6, <=3.9
1.7.1	0.8.2	>=3.6, <=3.9
1.7.0	0.8.1	>=3.6, <=3.8
1.7.0	0.8.0	>=3.6, <=3.8
1.6.0	0.7.0	>=3.6, <=3.8
1.5.1	0.6.1	>=3.5, <=3.8
1.5.0	0.6.0	>=3.5, <=3.8
1.4.0	0.5.0	==2.7, >=3.5, <=3.8
1.3.1	0.4.2	==2.7, >=3.5, <=3.7
1.3.0	0.4.1	==2.7, >=3.5, <=3.7
1.2.0	0.4.0	==2.7, >=3.5, <=3.7
1.1.0	0.3.0	==2.7, >=3.5, <=3.7
<=1.0.1	0.2.2	==2.7, >=3.5, <=3.7

 查找对应的关系  Previous PyTorch Versions | PyTorch

# CUDA 9.2
conda install pytorch==1.2.0 torchvision==0.4.0 cudatoolkit=9.2 -c pytorch

# CUDA 10.0
conda install pytorch==1.2.0 torchvision==0.4.0 cudatoolkit=10.0 -c pytorch

# CPU Only
conda install pytorch==1.2.0 torchvision==0.4.0 cpuonly -c pytorch

因为我的python版本为3.6.5，我最后选择的版本 cu101+torch1.7.1+torchvision0.8.2

下载torch-1.7.1+cu101和torchvision-0.8.2+cu101这2个文件安装，离线包下载地址 https://download.pytorch.org/whl/torch_stable.html

pip install F:\torch-1.7.1+cu101-cp36-cp36m-win_amd64.whl
pip install F:\torchvision-0.8.2+cu101-cp36-cp36m-win_amd64.whl

安装成功后，再执行训练脚本

python detect.py --source data\images\zidane.jpg --weights yolov5s.pt --conf 0.25

这时使用的就是gpu执行了。

训练指令

 python train.py --data fruit_training/dataset.yaml --cfg fruit_training/yolov5s.yaml --weights yolov5s.pt --batch-size 4 --epochs 100