使用自定义数据集在PyTorch中训练Yolo进行对象检测

在上一个故事中，我展示了如何使用预训练的Yolo网络进行物体检测和跟踪。 现在，我想向您展示如何使用由您自己的图像组成的自定义数据集重新训练Yolo。

挑战涉及检测隧道网络中的9个不同对象-它们是非常具体的对象，而不是标准Yolo模型中包含的常规对象。 在此示例中，我假设只有3个对象类。
有几种方法可以执行此操作，并且可以根据官方规范在训练脚本中定义图像，配置，批注和其他数据文件的位置，但这是一种更简单且井井有条的方法 Yolo的最佳做法。

Folder Structure

首先，您需要使用以下文件夹结构将所有训练图像放在一起（文件夹名称以斜体显示）：

Main Folder
--- data
    --- dataset name
        --- images
            --- img1.jpg
            --- img2.jpg
            ..........
        --- labels
            --- img1.txt
            --- img2.txt
            ..........
        --- train.txt
        --- val.txt

现在，让我们看看这些文件的外观（除了显而易见的图像文件）。
首先，注释文件。 每个图像需要一个.txt文件（相同名称，不同扩展名，单独的文件夹）。 每个文件仅包含一行，其格式为：

class x y width height

So, for example one file (for class 1) could be:

1 0.351466 0.427083 0.367168 0.570486

请注意，坐标和尺寸是整个图像尺寸的一部分。 例如，如果文件为600x600px，则坐标（200,300）将表示为（0.333333，0.5）。
train.txt和val.txt包含训练和验证图像的列表，每行一个，并带有完整路径。 例如，我的系统上此类文件的前两行是：

/datadrive/Alphapilot/data/alpha/images/IMG_6723.JPG
/datadrive/Alphapilot/data/alpha/images/IMG_6682.JPG

I used the following program to generate the 2 files, based on a 90% training / 10% validation split:

import glob
import os
import numpy as np
import sys
current_dir = "./data/artifacts/images"
split_pct = 10  # 10% validation set
file_train = open("data/artifacts/train.txt", "w")  
file_val = open("data/artifacts/val.txt", "w")  
counter = 1  
index_test = round(100 / split_pct)  
for fullpath in glob.iglob(os.path.join(current_dir, "*.JPG")):  
  title, ext = os.path.splitext(os.path.basename(fullpath))
  if counter == index_test:
    counter = 1
    file_val.write(current_dir + "/" + title + '.JPG' + "\n")
  else:
    file_train.write(current_dir + "/" + title + '.JPG' + "\n")
    counter = counter + 1
file_train.close()
file_val.close()

Creating Annotation Files

Now you are going to ask how to get the .txt annotation files. Well, I’m using a modified version of the BBOX tool, which is included in the Github repo. The way it works is this: You place the training images in different folders for each class. Look under the LoadDir function to figure out the folder structure (or modify for yours) — in my example, I have two folders, “forboxing” for images, and “newlabels” for the generated annotations, and under “forboxing” there are subfolders for each class (“0”, “1”, etc). You’ll have to modify the self.imgclass attribute at the top of the file, and run it separately for each class. This procedure makes everything a bit quicker. Using the tool itself is very intuitive — you just draw a square box around the object in each frame, then go to the next.

Config Files

Now for the config files in the config/ folder. First, coco.data would look like this:

classes = 3
train=data/alpha/train.txt
valid=data/alpha/val.txt
names=config/coco.names
backup=backup/

我认为这很不言自明。 backup参数未使用，但似乎是必需的。 coco.names 文件非常简单，它应该每行列出一个类的名称（对于注释文件，第一个对应于0，后跟1，依此类推）。 就我而言，该文件包含三个类：

DRILL
EXTINGUISHER
RANDY

Now, the most important of the configuration files is yolov3.cfg. It’s a big file, but here are the main things you have to change:
In the first [net] section, adjust the batch value and subdivisions to fit your GPU memory. The larger the batch size, the better and faster the training, but the more memory it will take. For an Nvidia GPU with 11Gb memory, a batch of 16 and 1 subdivision is good. Also here you can adjust the learning_rate.
Now, the most important (because if they’re not set correctly, your training program will fail) are the classes and final layer filters values. And you have to do it in three different places in the file. If you search the file, you’ll find 3 [yolo] sections. Inside that section, set classes to the number of classes in your model. You also have to change the filters value in the [convolutional] section right above [yolo]. That value is equal to:

filters = (classes + 5) x 3

So for my 3 classes, there are 24 filters. Be careful that this is only true for Yolo V3. V2 has a different formula.

Running the Training Script

现在您已经准备好进行实际的训练！ 训练程序是标准的Yolo脚本。 在“配置”部分中，设置所需的epoch数，确保文件夹路径正确，然后运行。 根据训练图像和硬件的数量，这可能需要几个小时到一天以上的时间。
该脚本将在每个时期后保存…抓取最后一个文件并将其放回config文件夹中，然后就可以对自定义数据集进行对象检测了！