MindSpore：YOLOv3人体目标检测模型实现（三）

4.模型训练

这一部分可以参见 附件\train_yolov3.py，官方平台有更通用的训练代码（参见：https://gitee.com/mindspore/models/tree/r1.5/official/cv/yolov3_darknet53）。

4.1 导包

首先是导入包：

import time

from mindspore import context
from mindspore.common import set_seed

from mindspore.train.serialization import load_checkpoint, load_param_into_net
from mindspore import Tensor
from mindspore.nn.optim.momentum import Momentum
from mindspore import amp
from mindspore.train.loss_scale_manager import FixedLossScaleManager
from mindspore.train.callback import ModelCheckpoint, RunContext
from mindspore.train.callback import CheckpointConfig

# 以下包请见附件代码 train_yolov3.py
from utils.utils_yolov3 import AverageMeter, get_param_groups, keep_loss_fp32, load_yolov3_backbone
from model.yolo import YOLOV3DarkNet53, YoloWithLossCell
from utils.kaiming_init import default_recurisive_init
from dataset.voc2012_dataset import create_voc2012_dataset
from utils.lr_generator import get_lr
from yolov3_train_conf import TrainYOLOv3Conf

set_seed(1)

上面的包可能有些没有用但我忘删了。

附件中引入的代码：

AverageMeter：主要是用于管理训练过程中的损失值，每个迭代一定次数给出这一段训练的损失平均值。

get_param_groups：主要是在配置优化器阶段使用。（keep_loss_fp32：我也不知道它是干什么的）

load_yolov3_backbone：主要是将Darknet53主干网络的预训练参数加载进来。

default_recurisive_init：使用He初始化，因为网络中大量使用Relu。

get_lr：学习率使用余弦退火函数生成，这个函数会根据参数生成每一代训练的学习率。

TrainYOLOv3Conf：训练配置。

4.2 设置context

然后设置Context：

# set context
context.set_context(mode=context.GRAPH_MODE,
                    device_target="GPU", save_graphs=False)
# set_graph_kernel_context
context.set_context(enable_graph_kernel=True)
context.set_context(graph_kernel_flags="--enable_parallel_fusion "
                                       "--enable_trans_op_optimize "
                                       "--disable_cluster_ops=ReduceMax,Reshape "
                                       "--enable_expand_ops=Conv2D")
# Set mempool block size for improving memory utilization, which will not take effect in GRAPH_MODE
if context.get_context("mode") == context.PYNATIVE_MODE:
    context.set_context(mempool_block_size="31GB")

这里面我认为用处比较大的是第一个set_context，尤其是device_target参数决定了运行环境是GPU。其它的参数我也不好评价它们的重要性（因为我也不太知道是啥意思），大家可以参考官网。

4.3 初始化配置和网络

接下来是初始化训练配置对象和网络：

# config
config = TrainYOLOv3Conf()

loss_meter = AverageMeter('loss')

# network
network = YOLOV3DarkNet53(is_training=config.data_training, config=config)
default_recurisive_init(network)
# ~ load_yolov3_learn(network, config.ckpt_path) # 用于迁移YOLOv3模型时使用
load_yolov3_backbone(network, config.backbone_ckpt_path)
network = YoloWithLossCell(network, config=config)

AverageMeter如前所述是用于管理训练过程中的损失值的。TrainYOLOv3Conf生成的对象config中保存有所有的模型训练配置，包括这里的网络配置、后面的读数据集配置以及训练配置，大家可以在 附件\yolov3_train_conf.py 中查看和编辑，TrainYOLOv3Conf的定义和其中各字段的含义如下：

"""训练YOLOv3时的参数配置"""
class TrainYOLOv3Conf():
    def __init__(self):
        
        # ---- dataset ----
        
        self.data_path = "VOCdevkit/VOC2012/"      # 数据集路径，要求到有Annotations和ImageSets以及JPEGImages的目录下
        self.data_usage = "my_person_train"        # ImageSets/Main/中的文件名，不包括txt后缀，说明读取哪些数据
        self.data_training = True                  # 数据用于训练，训练时将采用图像增强操作和shulffle
        
        self.num_classes = 1                       # 只检测一个类即person
        
        self.class_to_idx = {}
        
        self.anchor_scales = [[15, 38],            # 9个先验框大小
                              [34, 86],
                              [84, 127],
                                
                              [51, 192],
                              [91, 257],
                              [173, 195],
                                
                              [142, 319],
                              [221, 339],
                              [351, 365]]
        
        self.jitter = 0.3     # 0.3                # 图像缩放裁剪，见transform.py中的_choose_candidate_by_constraints函数
        self.scale_sma = 0.7  # 0.25
        self.scale_lar = 1.5  # 2
        
        self.hue = 0.1                             # 图像颜色随机变换。见transform.py中的color_distortion函数
        self.saturation = 1.5
        self.value = 1.5
        
        self.batch_size = 48                       # batch size
        self.max_box = 32                          # 训练时每张图像最多检测几个框
        
        self.label_smooth = 0                      # 用于损失函数中的一个设置
        self.label_smooth_factor = 0.1
        
        # ---- Model ----
        
        self.out_channel = 3 * ( 5 + self.num_classes) # 输出通道数
        
        self.keep_detect = True                        # 默认为True
        
        # ~ self.ckpt_path = "./yolo3v-280_36120.ckpt"         # 迁移训练模型
        self.backbone_ckpt_path = "./ckpt/darknet53_wp.ckpt"  # 加载预训练文件路径
        
        # ---- Train ----
        
        self.ignore_threshold = 0.7           # 用于训练时决定
        
        self.max_epoch = 320                  # 总训练的epoch
        self.pretrained_epoch_num = 0         # 已经训练的epoch
        self.steps_per_epoch = -1             # 每个epoch有多少个batch，这里在train_yolov3.py运行时再设置
        
        # lr, use in get_lr()
        self.lr = 0.001                       # 初始学习率
        self.warmup_epochs = 4
        self.T_max = 320
        self.eta_min = 0.0
        
        # opt: momentum
        self.momentum = 0.9        # use in Momentum init
        
        self.loss_scale = 1024     # use in Momentum and FixedLossScaleManager
        
        self.weight_decay =  0.016 # use in group_params
        
        # ckpt
        self.ckpt_interval = 10               # 隔多少epoch保存以下模型
        self.save_ckpt_path = "train_ckpt/"   # 保存epoch的路径
        
        # log
        self.log_interval = 50                # 训练多少step显示一下损失值和其它中间信息

网络初始化过程中，首先是生成了YOLOV3DarkNet53，然后对整个网络进行He初始化，再加载主干网络的参数进行覆盖（加载时会提示有些模块没加载，不要紧张，还记得结构图中主干网络后面的3条尾巴吗），最后再将它们放到YoloWithLossCell中。

一般来说，目标检测模型要使用迁移训练。本项目的YOLOv3模型其实基本上是copy的官网的（ModelZoo，https://gitee.com/mindspore/models/tree/r1.5/official/cv/yolov3_darknet53），如果大家能够找到别人用大数据集预训练好的模型，可以直接迁移。MindSpore的模型参数采用的是字典格式存储，键是字符串，其命名为你的模型class中的 字段.下一级字段...，若迁移时系统提示冲突，可以根据自己的模型字段重新建一个新字典加载。若找不到预训练的模型，在附件中我提供了利用Caltech-256分类数据集对主干网络进行预训练的方法和代码，本项目用着效果还行，要注意Caltech-256下好后内部其实有257个分类，要删掉最后一个。

4.4 加载数据集

# dataset
voc2012_dat, data_size = create_voc2012_dataset(config, 2)
config.steps_per_epoch = int(data_size / config.batch_size)
print("bath num in 1 epoch: ", config.steps_per_epoch)

4.5 生成Momentum优化器

# opt
lr = get_lr(config.lr,
            config.steps_per_epoch,
            config.warmup_epochs,
            config.max_epoch,
            config.T_max,
            config.eta_min,
            config.pretrained_epoch_num)
opt = Momentum(params=get_param_groups(network),
               learning_rate=Tensor(lr),
               momentum=config.momentum,
               weight_decay=config.weight_decay,
               loss_scale=config.loss_scale)
loss_scale_value = 1.0
loss_scale = FixedLossScaleManager(loss_scale_value, drop_overflow_update=False)
network = amp.build_train_network(network, optimizer=opt, loss_scale_manager=loss_scale,
                                  level="O2", keep_batchnorm_fp32=False)
keep_loss_fp32(network)

其中，get_lt 的参数分别是：

config.lr：设置初始学习率。

config.steps_per_epoch：一个epoch中训练多少steps，在数据集读取部分计算出该值。

warmup_epochs：代表学习率线性提升阶段（从config.eta_min到config.lr）的epoch数。

config.max_epoch：一共训练多少epoch。

config.T_max：余弦函数的周期是多少epoch。

config.eta_min：学习率最小值。

config.pretrained_epoch_num：已经训练的epoch数。

4.6 设置训练过程中保存模型

# ckpt
config.ckpt_interval = config.steps_per_epoch * config.ckpt_interval
print("config.ckpt_interval: ",config.ckpt_interval,config.steps_per_epoch)
# checkpoint save
ckpt_max_num = config.max_epoch * config.steps_per_epoch // config.ckpt_interval
ckpt_config = CheckpointConfig(save_checkpoint_steps=config.ckpt_interval,
                               keep_checkpoint_max=ckpt_max_num)
ckpt_cb = ModelCheckpoint(config=ckpt_config,
                          directory=config.save_ckpt_path,
                          prefix='yolov3')
cb_params = InternalCallbackParam()
cb_params.train_network = network
cb_params.epoch_num = ckpt_max_num
cb_params.cur_epoch_num = config.pretrained_epoch_num+1
run_context = RunContext(cb_params)
ckpt_cb.begin(run_context)

这里主要是设置保存模型的间隔（config.ckpt_interval 在配置中本来保存的是多少 epoch 保存一下模型文件，这里乘以每个epoch的训练代数，得出每隔多少代保存模型），以及保存模型的路径（config.save_ckpt_path）和保存的文件名前缀。这里 InternalCallbackParam 类的定义请见 附件\train_yolov3.py。

4.7 训练

# train
old_progress = -1
p_epoch = config.pretrained_epoch_num
p_step = p_epoch*config.steps_per_epoch
t_end = time.time()
data_loader = voc2012_dat.create_dict_iterator(output_numpy=True, num_epochs=1)

for i, data in enumerate(data_loader):
    images = data["image"]
    input_shape = images.shape[2:4]

    images = Tensor.from_numpy(images)

    batch_y_true_0 = Tensor.from_numpy(data['bbox1'])
    batch_y_true_1 = Tensor.from_numpy(data['bbox2'])
    batch_y_true_2 = Tensor.from_numpy(data['bbox3'])
    batch_gt_box0 = Tensor.from_numpy(data['gt_box1'])
    batch_gt_box1 = Tensor.from_numpy(data['gt_box2'])
    batch_gt_box2 = Tensor.from_numpy(data['gt_box3'])

    loss = network(images, batch_y_true_0, batch_y_true_1, batch_y_true_2, batch_gt_box0, batch_gt_box1,
                   batch_gt_box2)
    loss_meter.update(loss.asnumpy())

    # ckpt progress
    cb_params.cur_step_num = i + 1  # current step number
    cb_params.batch_num = i + 2
    ckpt_cb.step_end(run_context)

    if i % config.log_interval == 0:
        time_used = time.time() - t_end
        epoch = int(i / config.steps_per_epoch)
        per_step_time = time_used / config.log_interval * 1000
        fps = config.batch_size * (i - old_progress) / time_used
        print(
            'epoch[{}], iter[{}], {}, {:.2f} imgs/sec, lr:{},'
            ' per step time:{}ms'.format(epoch+p_epoch, i+p_step, loss_meter, fps, lr, per_step_time))
        t_end = time.time()
        loss_meter.reset()
        old_progress = i
    
    if (i + 1) % config.steps_per_epoch == 0:
        cb_params.cur_epoch_num += 1
        
print("Train Finish")

注意：训练过程中若超参数设置不好也可能导致模型训练效果不佳，例如训练的epoch设置得太小。

（未完，请见 YOLOv3人体目标检测模型实现（四））