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
    评论
您好!针对Mx-yolov3模型的训练和K210模型的使用,您可以按照以下步骤进行操作: 1. 准备数据集:首先,您需要准备一个包含标记好的目标边界框的数据集。确保每个图像文件都有对应的标签文件,标签文件的格式可以是YOLO格式或者其他常用格式,如Pascal VOC。 2. 配置Mx-yolov3模型:根据您的需求,您可以在MxNet中下载预训练yolov3模型,然后根据您的数据集和目标进行微调。您可以参考MxNet官方文档来了解如何配置和使用Mx-yolov3模型。 3. 数据增强和预处理:为了提高模型的泛化能力,您可以使用数据增强技术来扩充您的数据集。数据增强包括随机裁剪、缩放、旋转、翻转等操作。此外,还需要对输入图像进行预处理,例如归一化和调整大小。 4. 模型训练:使用准备好的数据集和配置好的模型,您可以开始训练过程了。根据您的硬件资源和实际情况,可以选择在本地或者云端进行训练。通常情况下,训练过程需要花费一定的时间和资源。 5. 模型转换和部署:在训练完成后,您需要将模型转换为适用于K210芯片的格式。K210芯片使用的是KPU神经网络加速器,所以您需要将模型转换为KPU可以识别的格式。您可以使用K210提供的工具和库来进行模型转换和部署。 请注意,以上步骤仅为一般指导,具体操作可能因您的环境和需求而有所不同。建议您参考相关文档和教程,以获得更详细的指导和支持。祝您成功训练和使用K210模型!如果您有其他问题,请随时提问。
评论 8
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值