当前位置: 代码迷 >> 综合 >> Windows10+pycharm+cuda11.0 yoloV4训练自己的数据集
  详细解决方案

Windows10+pycharm+cuda11.0 yoloV4训练自己的数据集

热度:4   发布时间:2024-01-21 02:09:23.0

yolov4的发布引起了不少的关注,但由于darknet是大佬c语言写的,并且windows环境下的编译确实很麻烦,所以找了一个不用编译的yolov4开源代码基于pytorch环境。

参考:https://blog.csdn.net/myr503270510/article/details/109642901

本人的项目环境是:Windows10+pycharm+cuda11.0

数据格式:coco数据集格式(数据标注利用百度飞桨)

下载yolov4_pytorch源码:https://github.com/Tianxiaomo/pytorch-YOLOv4

下载权重:yolov4.weights 权重已经github上给出 ,自行下载,下载完成后,在上面代码文件的目录下创建 weights 文件夹,将权重文件放入这个文件夹中

1.环境搭建

  • 创建环境

    进入conda命令行创建虚拟环境,环境名为torch36(环境名自己起,注意和其他环境区分开),python版本为3.6

      conda create --name yolov4 python=3.6

  • 激活环境

     conda activate yolov4

  • 安装pytorch

      需要根据自己电脑的cuda版本安装对应的pytorch版本,我的电脑cuda是11.0,因此 torch ==1.7.1,torchvision==0.8.2

     可以去pytorch官网查看版本

      conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c conda-forge

到这里,基础的环境已经配置好了,注意,上面配置的是 GPU版 的。

2.demo测试

环境配置好之后,可以跑一下demo文件测试一下

进入到下载好的代码文件夹下,执行如下命令,执行前确保相应的文件在对应文件夹下

python demo.py -cfgfile cfg/yolov4.cfg -weightfile weights/yolov4.weights -imgfile data/dog.jpg

3.训练自己的数据集

数据转换:

利用tool/coco_annotation.py将coco数据集的json文件转换成train.txt

生成的数据格式如下:

准备train.txt,内容是图片名和box。格式如下```
image_path1 x1,y1,x2,y2,id x1,y1,x2,y2,id x1,y1,x2,y2,id ...
image_path2 x1,y1,x2,y2,id x1,y1,x2,y2,id x1,y1,x2,y2,id ...
...
```
- image_path : 图片名
- x1,y1 : 左上角坐标
- x2,y2 : 右下角坐标
- id : 物体类别

模型训练:

参数修改,修改cfg.py文件

开始训练的时候我直接用原来的参数,batch size设为64,跑了几个epoch发现不对,我数据一共才二十多个。

后修改网络更新策略,不是按照每个epoch的step更新,使用总的steps更新,观察loss貌似可以训练了,后面打开电脑一看,loss收敛到2.e+4下不去了,此种必又蹊跷,遂kill了。

于是把batch size直接设为4,可以正常训练了。

Cfg.batch = 4

Cfg.subdivisions = 1 

开始训练

python train.py -l 0.001 -g 0 -pretrained yolov4.conv.137.pth -classes 3 -dir data/

-l 学习率

-g gpu id

-pretrained 预训练的主干网络,从AlexeyAB给的darknet的yolov4.conv.137转换过来的

-classes 类别种类 -dir 图片所在文件夹

看下loss曲线 

tensorboard --logdir log 

如果不能tensorboard,可以pip install tensorboard 然后再查看。

验证


#python model.py num_classes weightfile imagepath namefilepython model.py 3 weight/Yolov4_epoch166_coins.pth data/1.jpg data/data.names

 


 

 

 

 

 

 

  相关解决方案