【目标检测代码实战】从零开始动手实现yolov3：训练篇（一）

前言

 在前面几篇文章中小糖豆为大家讲解了yolo系列算法的演变。俗话说，光说不练假把式。接下来小糖豆将带领大家从零开始，亲自动手实现yolov3的训练与预测。

 本教程说明：

• 需要读者已经基本了解pytorch的基础用法以及yolov3的算法原理（不太了解的小伙伴可以移步yolo系列算法三）；
  

• 参考的代码：

  https://github.com/eriklindernoren/PyTorch-YOLOv3

• 小糖豆在原代码基础上修改了部分源码（接下来会按照修改后的源码进行讲解，主要修改点包括tensorboard可视化部分、voc数据集标签制作部分、模型参数配置等），同时增加了camera.py（调用usb相机实现yolov3实时检测，调用方式为：python3 camera.py），有需要的小伙伴可以关注公众号并后台回复：yolov3，免费获取源代码。

配置

 在开始讲解前，小糖豆建议大家先将源码下载下来配置好，先运行一下看看模型效果。在跟随小糖豆讲解的过程中每一步调试一下，多尝试多动手，可以起到事半功倍的奇效。

 我们先大致看下源码的组成部分：

• checkpoints：模型训练节点存放地址
  

• config：模型和数据集的配置文件
  

• data：训练所用的数据集

• output：图像测试输出地址
  

• utils：模型构建所需的模块（数据增强，数据解析，tensorboard可视化等）
  

• weights：预训练模型文件

• camera.py：模式实时预测
  

• detect.py：模型预测

• models.py：模型构建模块

• test.py：模型验证

• train.py：模型训练

 我们从train.py文件开始分析：

初始化

导入模块：

from __future__ import division
from models import *from utils.utils import *from utils.datasets import *from utils.parse_config import *from test import evaluatefrom terminaltables import AsciiTable
import osimport sysimport timeimport datetimeimport argparseimport cv2import torchimport torchvisionfrom torch.utils.data import DataLoaderfrom torchvision import datasetsfrom torchvision import transformsfrom torch.autograd import Variableimport torch.optim as optimfrom tensorboardX import SummaryWriter

命令行参数解析：

if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--epochs", type=int, default=100, help="number of epochs") parser.add_argument("--learning_rate", default=0.001, help="learning rate") parser.add_argument("--batch_size", type=int, default=12, help="size of each image batch") parser.add_argument("--gradient_accumulations", type=int, default=2, help="number of gradient accums before step") parser.add_argument("--model_def", type=str, default="config/yolov3-voc0712.cfg", help="path to model definition file") parser.add_argument("--data_config", type=str, default="config/voc0712.data", help="path to data config file") parser.add_argument("--pretrained_weights", type=str, default='./weights/darknet53.conv.74',help="if specified starts from checkpoint model") parser.add_argument("--n_cpu", type=int, default=1, help="number of cpu threads to use during batch generation") parser.add_argument("--img_size", type=int, default=416, help="size of each image dimension") parser.add_argument("--checkpoint_interval", type=int, default=1, help="interval between saving model weights") parser.add_argument("--evaluation_interval", type=int, default=1, help="interval evaluations on validation set") parser.add_argument("--compute_map", default=True, help="if True computes mAP every tenth batch") parser.add_argument("--multiscale_training", default=True, help="allow for multi-scale training") parser.add_argument("--conf_thres", type=float, default=0.8, help="object confidence threshold") parser.add_argument("--nms_thres", type=float, default=0.4, help="iou thresshold for non-maximum suppression") opt = parser.parse_args() print("type(opt) --> ",type(opt)) print(opt)

argparse是一个命令行参数解析模块，add_argument添加可能会发生变化的参数（比如模型的配置文件，训练批大小等），parse_args()函数将添加的参数进行解析。print打印结果如下：

type(opt) --> <class 'argparse.Namespace'>Namespace(batch_size=12, checkpoint_interval=1, compute_map=True, conf_thres=0.8, data_config='config/voc0712.data', epochs=100, evaluation_interval=1, gradient_accumulations=2, img_size=416, learning_rate=0.001, model_def='config/yolov3-voc0712.cfg', multiscale_training=True, n_cpu=1, nms_thres=0.4, pretrained_weights='./weights/darknet53.conv.74')

模型构建：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = Darknet(opt.model_def,img_size=opt.img_size).to(device)

使用Darknet类定义一个model 对象，Darknet类的定义位于model.py。在定义对象的时候会执行类的构造函数__init__。

class Darknet(nn.Module): def __init__(self, config_path, img_size=416): #config_path = "config/yolov3-voc0712.cfg" super(Darknet, self).__init__()  self.module_defs = parse_model_config(config_path) #解析模型参数 self.hyperparams, self.module_list = create_modules(self.module_defs) self.yolo_layers = [layer[0] for layer in self.module_list if hasattr(layer[0], "metrics")] self.img_size = img_size self.seen = 0 self.header_info = np.array([0, 0, 0, self.seen, 0], dtype=np.int32)

Darknet类继承自pytorch的nn.Module，parse_model_config函数以模型的配置文件为参数，返回的列表作为create_modules函数参数来构建pytorch模块。

我们先看看parse_model_config函数：

def parse_model_config(path): """Parses the yolo-v3 layer configuration file and returns module definitions""" file = open(path, 'r') lines = file.read().split('\n') lines = [x for x in lines if x and not x.startswith('#')] lines = [x.rstrip().lstrip() for x in lines] # get rid of fringe whitespaces

首先读取cfg文件，将内容保存在lines列表中，保存的时候去掉注释行、空行和无用空格。

module_defs = [] for line in lines: if line.startswith('['): # This marks the start of a new block module_defs.append({}) module_defs[-1]['type'] = line[1:-1].rstrip() if module_defs[-1]['type'] == 'convolutional': module_defs[-1]['batch_normalize'] = 0 else: key, value = line.split("=") value = value.strip()    module_defs[-1][key.rstrip()] = value.strip() return module_defs

遍历lines，将cfg文件中的每个块存储为dict。这些块的属性和值都以键值的形式存储在dict中。

create_modules函数使用返回的module_defs列表构建pytorch模块。

def create_modules(module_defs): """ Constructs module list of layer blocks from module configuration in module_defs """ hyperparams = module_defs.pop(0)

在迭代module_defs列表前，先获取模型的超参数。

output_filters = [int(hyperparams["channels"])] module_list = nn.ModuleList()

create_modules函数会返回超参数和module_list。module_list是一个nn.Sequential()，这个类等同于一个包含nn.Module对象的普通列表。

 定义一个卷积层必须定义它的卷积核维度。虽然卷积核的高度和宽度由 cfg 文件提供，但卷积核的深度是由上一层的卷积核数量（或特征图深度）决定的。这意味着我们需要持续追踪被应用卷积层的卷积核数量。

 路由层（route layer）从前面层得到特征图（可能是拼接的）。如果在路由层之后有一个卷积层，那么卷积核将被应用到前面层的特征图上，精确来说是路由层得到的特征图。因此，我们不仅需要追踪前一层的卷积核数量，还需要追踪之前每个层。

 为了实现上述过程，在迭代过程中我们将每个模块的输出卷积核数量添加到 output_filters列表中。接下来我们可以开始迭代module_defs列表了：

for module_i, module_def in enumerate(module_defs): modules = nn.Sequential()

nn.Sequential类被用于按顺序地执行 nn.Module 对象中的模型层。在 cfg 文件中一个模块可能包含多个层。例如，一个 convolutional 类型的模块有一个批量归一化层、一个 leaky ReLU 激活层以及一个卷积层。我们使用 nn.Sequential 将这些层串联起来，使用 add_module 函数创建单个模型块。以下展示了创建卷积层和上采样层的例子：

if module_def["type"] == "convolutional": bn = int(module_def["batch_normalize"]) filters = int(module_def["filters"]) kernel_size = int(module_def["size"]) pad = (kernel_size - 1) // 2 modules.add_module( f"conv_{module_i}", nn.Conv2d( in_channels=output_filters[-1], out_channels=filters, kernel_size=kernel_size, stride=int(module_def["stride"]), padding=pad, bias=not bn, ), ) if bn: modules.add_module(f"batch_norm_{module_i}", nn.BatchNorm2d(filters, momentum=0.9, eps=1e-5)) if module_def["activation"] == "leaky": modules.add_module(f"leaky_{module_i}", nn.LeakyReLU(0.1)) elif module_def["type"] == "upsample": upsample = Upsample(scale_factor=int(module_def["stride"]), mode="nearest") modules.add_module(f"upsample_{module_i}", upsample)

接下来我们看看创建路由层（route layer）和短接层（shortcut layer）：

elif module_def["type"] == "route": layers = [int(x) for x in module_def["layers"].split(",")] filters = sum([output_filters[1:][i] for i in layers]) modules.add_module(f"route_{module_i}", EmptyLayer())
elif module_def["type"] == "shortcut": filters = output_filters[1:][int(module_def["from"])] modules.add_module(f"shortcut_{module_i}", EmptyLayer())

创建路由层时，首先提取出关于层属性的值，保存在layers列表中。每一次访问到路由层时，它的layers参数可能有不止一个值。

• 当只有一个值时，它输出这一层通过该值索引的特征图。比如说-4意味着输出路由层往前第四个特征图。

• 当有两个值时，它将返回这两个值索引的拼接后的特征图。比如-1和61意味着输出前一层（-1）和第61层的特征图，并将它们按深度拼接。

由于路由层仅仅涉及到特征图的拼接，没有多余的操作，因此在创建模块时使用一个不执行任何操作的模型层，在执行前向传播时完成拼接即可。

同样地，创建短接层时仅仅涉及到特征图的相加，也采用空白模型层的方式。

所谓的空白模型层如下，什么操作也没有：

class EmptyLayer(nn.Module): """Placeholder for 'route' and 'shortcut' layers""" def __init__(self): super(EmptyLayer, self).__init__()

最后一个模块是yolo层，这一层保存了模型中的一些重要信息，如anchor box尺寸，类别数量，图像尺寸。在这里我们定义了一个YOLOLayer模型层，该模型层的具体细节等到前向传播的时候在介绍，这里小伙伴们只需要默认它是一个模型层就可以了。

elif module_def["type"] == "yolo": anchor_idxs = [int(x) for x in module_def["mask"].split(",")] # Extract anchors anchors = [int(x) for x in module_def["anchors"].split(",")] anchors = [(anchors[i], anchors[i + 1]) for i in range(0, len(anchors), 2)] anchors = [anchors[i] for i in anchor_idxs] num_classes = int(module_def["classes"]) img_size = int(hyperparams["height"]) # Define detection layer yolo_layer = YOLOLayer(anchors, num_classes, img_size) modules.add_module(f"yolo_{module_i}", yolo_layer) # Register module list and number of output filters module_list.append(modules) output_filters.append(filters)return hyperparams, module_list

想观察下模型架构的小伙伴可以运行下models.py。小糖豆在文件最后加了几行代码，可以打印构建的模型块。

if __name__ == "__main__": module_list = create_modules(parse_model_config("./config/yolov3-voc0712.cfg")) print(module_list)

会打印出模型信息，包含106个模型块，如下：

model.apply(weights_init_normal)

构建好模型后，对模型的参数初始化。

def weights_init_normal(m): classname = m.__class__.__name__ if classname.find("Conv") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) elif classname.find("BatchNorm2d") != -1: torch.nn.init.normal_(m.weight.data, 1.0, 0.02) torch.nn.init.constant_(m.bias.data, 0.0)

接下来加载预训练模型：

if opt.pretrained_weights: if opt.pretrained_weights.endswith(".pth"): model.load_state_dict(torch.load(opt.pretrained_weights)) else: model.load_darknet_weights(opt.pretrained_weights)

第一次训练时，一般会使用model.load_darknet_weights加载darknet预训练模型进行finetuning。

除此之外，训练过程可能会分好几个阶段，也可能会由于空间不足而意外中止，此时可以执行model.load_state_dict加载训练中保存的中间节点的模型文件。

for i,p in enumerate(model.named_parameters()): if i == 156: #darknet部分参数不更新 break p[1].requires_grad = False

如果不想改变yolov3的backbone基础网络部分的参数，可以通过requires_grad来设置。

到这一步模型已经搭建完毕，小糖豆将在下篇文章中讲解数据集的载入与预处理。

未完待续......