大名鼎鼎的YOLOv3由于其速度快，精度高等特点，成为了目前学术界以及工业界最广为使用的目标检测模型之一。原始的YOLOv3是基于darknet框架下完成了，因此本文将详细讲述如何在Pytorch下使用YOLOv3模型训练自己的数据集并进行预测。
本文将分为四部分进行讲述：
环境配置
数据集准备
模型训练
模型预测
1.环境配置
1.1 Python环境安装
Python的编译环境非常之多，小伙伴们可以根据自己喜好来安装Anaconda，Pycharm等等，在此不做赘述，本文使用的是Pycharm社区版。
1.2Pytorch安装
首先进入以下网址：/
在这里选择自己的系统，安装方式（conda安装，pip安装等等），语言（本文是Python），以及你的CUDA版本，就可以获得安装指令了。复制该条指令，打开CMD就可以进行pip安装或者conda安装了。
Tips1：如何查看自己的CUDA版本呢？
CUDA是Nvidia显卡特有的运算平台，能够在训练中使用GPU进行加速。如果你没有Nvidia独显或者只有AMD显卡，那就只能在上面的选项中选择“None”来安装CPU版本了。如果你拥有Nvidia独显，桌面右击，选择“Nvidia控制面板”，点击帮助，点击系统信息，点击组件，就可以看到你的CUDA版本啦，如下面三张图片所示。
Tips2：Pytorch安装太慢怎么办？
由于是国外的网站，经常会出现被墙的现象导致安装速度过慢甚至速度为0的现象。这里推荐大家使用清华镜像来进行安装，速度很快。方法如下。
打开CMD，逐步输入以下命令
conda config --add channels https://mirrors.tuna./anaconda/pkgs/free/conda config --add channels https://mirrors.tuna./anaconda/pkgs/main/conda config --set show_channel_urls yesconda config --add channels https://mirrors.tuna./anaconda/cloud/pytorch/
然后键入你的安装指令，比如：
conda install pytorch torchvision cudatoolkit=10.1
2. 数据集准备
首先，我们需要下载代码：
/ultralytics/yolov3
其次，准备数据集，本文使用的是Pascal Voc的数据集，总共约5000张图片。
数据集下载地址：
/projects/pascal-voc-dataset-mirror/
先看下该数据集中文件结构以及内容：
在这里我们只需要用到JPEGImages和Annotations这两个文件夹。前者是图片，jpg格式的。后者是每张图片对应的标签信息，xml格式的。标签信息中最主要的内容就是相对应的图片中包含的物体类别以及它们的坐标位置信息。将这两个文件夹放到项目中data文件夹下。
此时，我们需要对图片和标签进行预处理，需要在项目根目录下创建两个脚本文件:
(1).maketxt.py:
该脚本会在data/ImageSets文件夹下生成：trainval.txt, test.txt, train.txt, val.txt，内容如下图所示，将图片分成了训练集，验证集和测试集，并将文件名（不带扩展名）汇总在txt文件中。
#-*- coding:utf-8 -*import osimport randomtrainval_percent = 0.1train_percent = 0.9xmlfilepath = 'data/Annotations' txtsavepath = 'data/ImageSets' total_xml = os.listdir(xmlfilepath)num = len(total_xml)list = range(num)tv = int(num * trainval_percent)tr = int(tv * train_percent)trainval = random.sample(list, tv)train = random.sample(trainval, tr)ftrainval = open('data/ImageSets/trainval.txt', 'w') ftest = open('data/ImageSets/test.txt', 'w') ftrain = open('data/ImageSets/train.txt', 'w') fval = open('data/ImageSets/val.txt', 'w')for i in list:name = total_xml[i][:-4] + '\n'if i in trainval:ftrainval.write(name)if i in train:ftest.write(name)else:fval.write(name)else:ftrain.write(name)ftrainval.close()ftrain.close()fval.close()ftest.close()
(2). voc_label.py：
此脚本将在data文件夹下生成images和labels两个文件夹用于训练。images文件夹下储存训练图片，labels文件夹下储存相应的标签文件。这里的标签文件全部是由原xml文件转化而来的txt文件，且将坐标位置信息全部归一化处理，代码如下。
需要注意的是，代码中第9行的classes对应的是训练集中的标签，如果你使用本文中的训练集，那就无需改动；如果你使用自己的训练集，这里需要修改成你自己的标签。
#-*- coding:utf-8 -*import xml.etree.ElementTree as ETimport pickleimport osfrom os import listdir, getcwdfrom os.path import joinsets = ['train', 'test','val']classes = ['person','bird','cat','cow','dog','horse','bicycle','boat','bus','car','motorbike','train','bottle','chair','diningtable','pottedplant','sofa','tvmonitor']def convert(size, box):dw = 1. / size[0]dh = 1. / size[1]x = (box[0] + box[1]) / 2.0y = (box[2] + box[3]) / 2.0w = box[1] - box[0]h = box[3] - box[2]x = x * dww = w * dwy = y * dhh = h * dhreturn (x, y, w, h)def convert_annotation(image_id):in_file = open('data/Annotations/%s.xml' % (image_id))out_file = open('data/labels/%s.txt' % (image_id), 'w')tree = ET.parse(in_file)root = tree.getroot()size = root.find('size')w = int(size.find('width').text)h = int(size.find('height').text)for obj in root.iter('object'):difficult = obj.find('difficult').textcls = obj.find('name').textif cls not in classes or int(difficult) == 1:continuecls_id = classes.index(cls)xmlbox = obj.find('bndbox')b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),float(xmlbox.find('ymax').text))bb = convert((w, h), b)out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')wd = getcwd()print(wd)for image_set in sets:if not os.path.exists('data/labels/'):os.makedirs('data/labels/')image_ids = open('data/ImageSets/%s.txt' % (image_set)).read().strip().split()list_file = open('data/%s.txt' % (image_set), 'w')for image_id in image_ids:list_file.write('data/images/%s.jpg\n' % (image_id))convert_annotation(image_id)list_file.close()
执行完上述两个脚本即可开始进行训练。
3.开始训练
3.1 参数修改
训练之前，需要修改部分参数。
下图是train.py的参数部分。
首先，如果根据你GPU的显存大小，需要修改batchsize和训练图片的输入尺寸，如果过大，会报“CUDA out of memory”的错误,即显存溢出。这里的batchsize大小为batch-size*accumulate，本文中是8*1=8。图片尺寸是img-size(192*192)。红色框注出。
网络结构yolov3-spp.cfg文件以及yolo.data, yolo.names需要准备好，下文将详细叙述如何制作。黄色框注出。
其余参数可根据实际情况自行更改，非必需。
如果你有多个GPU，可以选择某单个GPU或者多个GPU同时参与训练，白色框注出，修改GPU id即可。
Tips3：如何查看GPU的显存占用情况？
为何确定最合适的batchsize和img_size，我们需要实时监控现存占用情况。
打开cmd，然后：
cd C:\Program Files\NVIDIA Corporation\NVSMInvidia-smi
如下所示：
同样，你可以使用nvidia-smi -l 1的命令，可以每秒刷新一次来进行监控。
3.2训练文件准备
yolo.names:
内容如下，这个储存训练集中所有的标签，如果你使用自己的训练集，就改成你自己的标签，储存在data文件夹下。
person,bird,cat,cow,dog,horse,bicycle,boat,bus,car,motorbike,train,bottle,chair,diningtable,pottedplant,sofa,tvmonitor
yolo.data：
内容如下，储存在data文件夹下。
classes=20train=data/train.txtvalid=data/val.txtnames=data/yolo.namesbackup=backup/
yolov3-spp.cfg:
cfg文件储存着yolo模型的网络结构。yolov3共有3个yolo层，因此在网络结构中，每个yolo层的上一层卷积层都需要修改filters的值，该值的计算公式为：
（classes + 5）* anchors_of_this_yolo
其中classes为训练集的标签数，本文中为20；5代表4个坐标位置+1个置信度；anchors_of_this_yolo代表该层yolo的anchor数目，本文总共有9个anchor，每层yolo用3个，所以这里是3。因此最终的filters = (20+5)*3 = 75。
这里务必三处都要修改，否则会报“shape xx is invalid for input of size xx”的错误。
最后需要注意的是：
anchors的大小对于最终的训练效果有着非常重要的影响，它和训练集有着非常密切的关系。这里我们姑且先使用默认的anchors值，下一期文章我将具体介绍如何计算训练集的anchors。
最后，运行train.py即可开始训练，训练后的模型将保存在weights文件夹中。
4.检验效果：
将待检测的图片放入data\sample文件夹下，基于训练获得的weights文件夹下的best.pt模型文件，运行detect.py即可完成检测，检测效果如下
总结
总体来说，训练与预测过程非常简单，最终的模型效果也是非常理想。但实际上，spp网络结构是yolo中网络结构最为复杂，算量最大的模型(141.45 Bn, 20FPS)，用于实验尚可，但基本无法用于工程实践之中，后续博主会不定期推出相关的剪枝算法等模型优化方法来提升模型的性能，欢迎继续关注。
最后，欢迎大家关注我的个人公众号，我将会不定期推出计算机视觉领域前沿研究和实用技术，一起交流与进步！
