介绍
gluon把mxnet再进行封装,封装的风格非常接近pytorch
使用gluon的好处是非常容易把pytorch模型向mxnet转化
唯一的问题是gluon封装还不成熟,封装好的layer不多,很多常用的layer 如concat,upsampling等layer都没有
这里关注如何把pytorch 模型快速转换成 mxnet基于symbol 和 exector设计的网络
pytorch转mxnet module
关键点:
- mxnet 设计网络时symbol 名称要和pytorch初始化中各网络层名称对应
- torch.load()读入pytorch模型checkpoint 字典,取当中的'state_dict'元素,也是一个字典
- pytorch state_dict 字典中key是网络层参数的名称,val是参数ndarray
- pytorch 的参数名称的组织形式和mxnet一样,但是连接符号不同,pytorch是'.',而mxnet是'_'比如:
pytorch '0.conv1.0.weight'
mxnet '0_conv1_0_weight'
- pytorch 的参数array 和mxnet 的参数array 完全一样,只要名称对上,直接赋值即可初始化mxnet模型
需要做的有以下几点:
- 设计和pytorch网络对应的mxnet网络
- 加载pytorch checkpoint
- 调整pytorch checkpoint state_dict 的key名称和mxnet命名格式一致
FlowNet2S PytorchToMxnet
pytorch flownet2S 的checkpoint 可以在github上搜到
import mxnet as mx
from symbol_util import *
import pickle
def get_loss(data, label, loss_scale, name, get_input=False, is_sparse = False, type='stereo'):
if type == 'stereo':
data = mx.sym.Activation(data=data, act_type='relu',name=name+'relu')
# loss
if is_sparse:
loss =mx.symbol.Custom(data=data, label=label, name=name, loss_scale= loss_scale, is_l1=True,
op_type='SparseRegressionLoss')
else:
loss = mx.sym.MAERegressionOutput(data=data, label=label, name=name, grad_scale=loss_scale)
return (loss,data) if get_input else loss
def flownet_s(loss_scale, is_sparse=False, name=''):
img1 = mx.symbol.Variable('img1')
img2 = mx.symbol.Variable('img2')
data = mx.symbol.concat(img1,img2,dim=1)
labels = {'loss{}'.format(i): mx.sym.Variable('loss{}_label'.format(i)) for i in range(0, 7)}
# print('labels: ',labels)
prediction = {}# a dict for loss collection
loss = []#a list
#normalize
data = (data-125)/255
# extract featrue
conv1 = mx.sym.Convolution(data, pad=(3, 3), kernel=(7, 7), stride=(2, 2), num_filter=64, name=name + 'conv1_0')
conv1 = mx.sym.LeakyReLU(data=conv1, act_type='leaky', slope=0.1)