Yolov3(Mxnet)FP16训练

在某些硬件下，FP16比FP32在可接受的精度损失下，训练、测试的加速效果明显。我们根据Mxnet中的Gluoncv，得到支持FP16的Yolov3模型。首先需要下载Gluoncv源码并将其修改，然后可以在本地训练中import更改的模型实现训练。

Gluoncv代码链接：https://github.com/dmlc/gluon-cv

实际上，Mxnet提供FP16和FP32网络模型转换的功能，如果只是想测试加速，在通过get_model或者load_parameters获得网络模型后，通过如下便可以将float32的网络整体转换成float16。

net.cast('float16')

而相应的数据在传进网络前也应该进行转换：

data = data.astype('float16')

但是如果在训练时使用FP16，由于有些网络定义中的参数写死，或者有些操作不支持FP16，因此直接在训练中使用net.cast()有时会报错，我们以Gluoncv中的Yolov3模型为例，并且使用原版的train_yolo3.py训练程序，实现FP16训练。

train_yolo3.py下载链接：https://gluon-cv.mxnet.io/model_zoo/detection.html#yolo-v3

1. 修改train_yolo3.py中的网络类型

原始代码中，有net和async_net两个网络，async_net注释写着“used by cpu worker”：

if args.syncbn and len(ctx) > 1:
    net = get_model(net_name, pretrained_base=True, norm_layer=gluon.contrib.nn.SyncBatchNorm, norm_kwargs={'num_devices': len(ctx)})
    async_net = get_model(net_name, pretrained_base=False)  # used by cpu worker
else:
    net = get_model(net_name, pretrained_base=True)
    async_net = net

按照注释说明，我们需要将net转为float16类型，而async_net不动：

train_data, val_data = get_dataloader(
        async_net, train_dataset, val_dataset, args.data_shape, args.batch_size, args.num_workers, args)
net.cast('float16')

2.  修改train_yolo3.py中训练和验证数据

在train函数中使用astype将训练和验证数据修改为float16格式：

obj_loss, center_loss, scale_loss, cls_loss = net(x.astype('float16', copy=False), gt_boxes[ix], *[ft[ix] for ft in fixed_targets])

同样validate函数中也做修改：

ids, scores, bboxes = net(x.astype('float16', copy=False))

3. 修改train_yolo3.py中的trainer，增加 multi_precision

trainer = gluon.Trainer(
        net.collect_params(), 'sgd',
        {'wd': args.wd, 'momentum': args.momentum, 'lr_scheduler': lr_scheduler, 'multi_precision': True},
        kvstore='local')

4. 修改gluoncv/model_zoo/yolo/yolo3.py中网络

因为计算loss、anchor iou等操作不支持FP16，所以在output层输出后，将输出立刻转成FP32可以避免很多麻烦：

将输出层原始代码：

pred = self.prediction(x).reshape((0, self._num_anchors * self._num_pred, -1))

增加转换类型的操作：

pred = self.prediction(x).astype('float32', copy=False)
pred = pred.reshape((0, self._num_anchors * self._num_pred, -1))

将输出层解码box_center和scale的部分增加类型转换，原始代码：

box_centers = F.broadcast_add(F.sigmoid(raw_box_centers), offsets) * self._stride
box_scales = F.broadcast_mul(F.exp(raw_box_scales), anchors)

修改为：

box_centers = F.broadcast_add(F.sigmoid(raw_box_centers), offsets.astype('float32', copy=False)) * self._stride
box_scales = F.broadcast_mul(F.exp(raw_box_scales), anchors.astype('float32', copy=False))

至此，训练程序修改完成，可以直接将修改的gluoncv包重命名为gluoncv_fp16,然后放到train_yolo3.py路径下，import里面的get_model，便可以训练FP16的Yolov3，不过训练的时候不能使用--syncbn，在Tesla V100上测试训练速度确实提升明显。如果训练的时候出现Nan，可以现在FP32下warmup，然后再使用FP16恢复训练，这样比较稳定。

from gluoncv_fp16.model_zoo import get_model

注：Mxnet已经支持Automatic Mixed Precision，看着简单高效，后续会尝试一下：

https://mxnet.incubator.apache.org/api/python/docs/tutorials/performance/backend/amp.html#Data-loader-and-helper-functions

而且，Gluoncv中的Faster rcnn训练程序train_faster_rcnn.py已经应用Automatic Mixed Precision：

https://gluon-cv.mxnet.io/model_zoo/detection.html#faster-rcnn

尝试了一下最近Yolo支持AMP的训练文件，但是发现不是很好用，一些操作在FP16下比较容易NAN，发现用AMP也是一样NAN。