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。

 

  • 1
    点赞
  • 4
    收藏
    觉得还不错? 一键收藏
  • 8
    评论
评论 8
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值