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也成功转换啦~

---

总结

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