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

4.模型训练
这一部分可以参见 附件\train_yolov3.py，官方平台有更通用的训练代码（参见：https://gitee.com/mindspore/m...）。

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/m...），如果大家能够找到别人用大数据集预训练好的模型，可以直接迁移。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人体目标检测模型实现（四））