使用 MMdnn IR 接口模型将 Pytorch 模型转换为 Keras 模型

Question

我从link 下载了一个预训练的 Pytorch cifar 模型，它具有自定义架构。我想将此模型转换为 Keras 模型。我知道有一些工具可以做到这一点。例如，我可以将 Pytorch 模型转换为标准模型，如 onnx 或 IR。

现在我已经使用IR接口模型和下面的代码来加载和保存整个模型：

import torch
import torch.nn as nn
import numpy as np
from  torch.autograd import Variable
import torch as th
from collections import OrderedDict

class CIFAR(nn.Module):
    def __init__(self, features, n_channel, num_classes):
        super(CIFAR, self).__init__()
        assert isinstance(features, nn.Sequential), type(features)
        self.features = features
        self.classifier = nn.Sequential(
            nn.Linear(n_channel, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

def make_layers(cfg, batch_norm=False):
    layers = []
    in_channels = 3
    for i, v in enumerate(cfg):
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            padding = v[1] if isinstance(v, tuple) else 1
            out_channels = v[0] if isinstance(v, tuple) else v
            conv2d = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=padding)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(out_channels, affine=False), nn.ReLU()]
            else:
                layers += [conv2d, nn.ReLU()]
            in_channels = out_channels
    return nn.Sequential(*layers)

n_channel=128
cfg = [n_channel, n_channel, 'M', 2*n_channel, 2*n_channel, 'M', 4*n_channel, 4*n_channel, 'M', (8*n_channel, 0), 'M']
layers = make_layers(cfg, batch_norm=True)
model = CIFAR(layers, n_channel=8*n_channel, num_classes=10)
pretrained=True

if pretrained:
    m = th.load('MY_PATH/cifar10-d875770b.pth')
    state_dict = m.state_dict() if isinstance(m, nn.Module) else m
    assert isinstance(state_dict, (dict, OrderedDict)), type(state_dict)
    model.load_state_dict(state_dict)

torch.save(model, 'MY_PATH/pytorch.pth')

现在是时候将上述 Pytorch 模型 pytorch.pth 转换为 IR 模型了。

在 cmd 中我输入了 here 示例中的这个命令：

mmtoir -f pytorch -d IRModel --inputShape 3,32,32 -n pytorch.pth

但是出现了这个错误：

Traceback (most recent call last):
  File "c:\users\***\anaconda3\lib\site-packages\mmdnn\conversion\pytorch\pytorch_parser.py", line 76, in __init__
    model = torch.load(model_file_name)
  File "c:\users\***\anaconda3\lib\site-packages\torch\serialization.py", line 387, in load
    return _load(f, map_location, pickle_module, **pickle_load_args)
  File "c:\users\***\anaconda3\lib\site-packages\torch\serialization.py", line 574, in _load
    result = unpickler.load()
AttributeError: Can't get attribute 'CIFAR' on <module '__main__' from 'C:\\Users\\***\\Anaconda3\\Scripts\\mmtoir.exe\\__main__.py'>

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "c:\users\***\anaconda3\lib\runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "c:\users\***\anaconda3\lib\runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "C:\Users\***\Anaconda3\Scripts\mmtoir.exe\__main__.py", line 9, in <module>
  File "c:\users\***\anaconda3\lib\site-packages\mmdnn\conversion\_script\convertToIR.py", line 192, in _main
    ret = _convert(args)
  File "c:\users\***\anaconda3\lib\site-packages\mmdnn\conversion\_script\convertToIR.py", line 92, in _convert
    parser = PytorchParser(model, inputshape[0])
  File "c:\users\***\anaconda3\lib\site-packages\mmdnn\conversion\pytorch\pytorch_parser.py", line 78, in __init__
    model = torch.load(model_file_name, map_location='cpu')
  File "c:\users\***\anaconda3\lib\site-packages\torch\serialization.py", line 387, in load
    return _load(f, map_location, pickle_module, **pickle_load_args)
  File "c:\users\***\anaconda3\lib\site-packages\torch\serialization.py", line 574, in _load
    result = unpickler.load()
AttributeError: Can't get attribute 'CIFAR' on <module '__main__' from 'C:\\Users\\***\\Anaconda3\\Scripts\\mmtoir.exe\\__main__.py'>

我该如何解决？

Answer 1

尽管这个问题已经很久没有提出来了，我还是会分享答案。

我的解决方案相当简单，我在 anaconda 环境中安装了 mmdnn。并且文件mmtoir 存储在以下位置：

'user_name/anaconda3/envs/env_name/bin/mmtoir'

并且需要在此文件中明确引用我的神经网络类。就您而言，这意味着您需要：

1) 找到此文件

2) 复制CIFAR类的定义

class CIFAR(nn.Module):
def __init__(self, features, n_channel, num_classes):
    super(CIFAR, self).__init__()
    assert isinstance(features, nn.Sequential), type(features)
    self.features = features
    self.classifier = nn.Sequential(
        nn.Linear(n_channel, num_classes)
    )

def forward(self, x):
    x = self.features(x)
    x = x.view(x.size(0), -1)
    x = self.classifier(x)
    return x

进入mmtoir 文件。保存并再次运行命令。

以上是解决这个问题的两种方法之一，比较简单。您可以阅读此discussion 了解更多详情。

注意： MMdnn 目前仅支持PyTorch=0.4.0。因此，请确保您的模型也接受了相同版本的训练。使用PyTorch>0.4.0，所有步骤都可以正常运行，但最后会抛出错误。