MMdnn:微软模型转换工具

Posted Yan_Joy

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了MMdnn:微软模型转换工具相关的知识,希望对你有一定的参考价值。

简介

深度学习火了,一同热的还有深度学习框架。当然经过了几年的淘汰(比如theano),目前最火的大概有Tensorflow、PyTorch以及Caffe了。他们各有所长,开发时可能会选定一个单独的平台,但是在应用时就需要考虑其他的因素。重新在其他框架上改代码就变得比较麻烦,如果能把模型直接转换好,更重要的是不同平台效果相同,那就会节省很多时间。MMdnn就是这样的一个模型转换工具,而且是微软出品,质量有保障。

下图是它转换的过程,左边是输入的模型框架,会被统一转换成IR文件,然后再输出成目标框架模型。


安装

pip install mmdnn

安装实际上很简单,但有一点注意的是它所支持的各平台版本,在提供的环境文件中(MMdnn/requirements/linux-py3.txt):

six==1.10.0
numpy==1.15.4
protobuf==3.6.1
pillow == 5.2.0
setuptools==39.1.0
tensorflow==1.13.1
keras==2.2.4
coremltools==2.1.0
mxnet==1.2.0
cntk==2.6
http://download.pytorch.org/whl/cpu/torch-0.4.0-cp35-cp35m-linux_x86_64.whl
https://raw.githubusercontent.com/PaulaQin/pycaffe_packages/master/caffe-1.0-py3-none-any.whl
torchvision==0.2.1
onnx==1.4.1
onnx-tf==1.2.1

可以看到,很多都不是最新框架的版本,所以对于很新的模型可能还会有一些问题。


使用

我的目标是将PyTorch模型转换为Caffe和Tensorflow模型。首先注意的一点是原始模型是用torch.save(model, filename)直接保存的,而不是只保存了'state_dict'。否则可能会出现:"AttributeError: 'collections.OrderedDict' object has no attribute 'state_dict'"这样的错误。

Caffe

我们使用官方提供的resnet18模型,通过mmdownload可以直接下载:

mmdownload -f pytorch -n resnet101 -o ./

然后进行转换:

mmconvert -sf pytorch -iw imagenet_resnet18.pth -df caffe -om caffe_resnet18_official --inputShape 3,224,224

这里注意的是要输入--inputShape
然后会输出caffe_resnet18_official.caffemodel参数文件和caffe_resnet18_official.prototxt模型文件。如果直接使用Caffe推荐的LMDB格式作为输入运行,就会得到很惨的结果(大概是0%)。这是转换出问题了吗?如果读取模型文件,会发现对应层的参数并没有什么问题。困扰了两天左右,终于发现了原因:想要得到相同的结果,要有和原模型一样的输入。
这是什么意思呢?模型转换只能保证运算参数相同,但测试时的数据输入与原模型训练时所经过的预处理函数不同,就会出现输入不一致的情况,这时输出肯定也不一致了。
PyTorch官方训练代码中给出了预处理过程:

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])


data_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder(args.data, transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            normalize,
        ])),
        batch_size=args.batch_size, shuffle=False,
        num_workers=2, pin_memory=True)

因此输入时也要按照这样的形式。由于命令行接口难以加入预处理函数,因此需要python实现:

model_converted = caffe.Net(args.ModelPath, args.WeightPath, caffe.TEST)
for i, (data, target) in enumerate(data_loader):
    input_data = data.numpy()
    target = target.numpy()
    model_converted.blobs[model_converted.inputs[0]].data[...] = input_data
    predict = model_converted.forward()[model_converted.outputs[-1]]
    converted_predict = np.squeeze(predict)
    ...

这是加入到caffe的预处理部分。再次运行,准确率和原来是一致的。


Tensorflow

mmconvert -sf pytorch -iw imagenet_resnet18.pth -df tensorflow -om tf_resnet18_official --inputShape 3,224,224

转换命令基本上没有区别,运行后输出一个目录。然后就是怎么用的问题了,官网上给出的是:

export_dir = "./tf_resnet152"
with tf.Session(graph=tf.Graph()) as sess:
    tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.TRAINING], export_dir)

    x = sess.graph.get_tensor_by_name('input:0')
    y = sess.graph.get_tensor_by_name('xxxxxx:0')
    ......
    _y = sess.run(y, feed_dict=x: _x)

有两点需要注意:

  1. tf.saved_model.tag_constants这一项,是与--dump_tag SERVING,TRAINING所对应的;
  2. 模型的输出是什么?这个东西不去看Graph是不清楚的,但这样就会需要读出来模型,再通过summary保存图,最后用tensorboard显示出来。一般来说,最后一层是全连接,如resnet18,输出就是dense/BiasAdd:0。还有一个方法可以手动改名,但就不能一步转换了,只能改IR模型文件最后一层的name(默认是空)。
    同样给出代码:
load = tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], args.model)
input_tf = sess.graph.get_tensor_by_name('input:0')
model_tf = sess.graph.get_tensor_by_name(args.tf_output_node)
init = tf.global_variables_initializer()
sess.run(init)

for i, (data, target) in enumerate(val_loader):
    input_data = data.numpy()
    input_data = np.transpose(input_data, (0, 2, 3, 1))
    target = target.numpy()
    predict = sess.run(model_tf, feed_dict = input_tf : input_data)
    converted_predict = np.squeeze(predict)
    ...

这和上面caffe还有一点不同,就是多了一次np.transpose维度变换,是一个小坑。
这样一来,tensorflow也成功转换啦~


总结

上面的坑都是自己踩的,网上似乎没有找到类似的教程,当时真的有点无语。。记录一下,希望能帮到别人吧。

以上是关于MMdnn:微软模型转换工具的主要内容,如果未能解决你的问题,请参考以下文章

能在不同的深度学习框架之间转换模型?微软的MMdnn做到了

微软开源 MMdnn:深度学习框架随心切换

陈程:MMdnn如何让模型在不同深度学习框架间转换?

MSRA 系统研究员陈程:MMdnn 如何让模型在不同深度学习框架间转换?

如何把Maya2012中制作的模型导入到微软的WPF中并且显示出来?

MxNet 模型转Tensorflow pb模型