在某些硬件下,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,看着简单高效,后续会尝试一下:
而且,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。