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

时间:2019-06-17 18:10:05

标签: python keras pytorch

我已从此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'>

我该如何解决?

1 个答案:

答案 0 :(得分:0)

尽管问这个问题已经很久了,但我将分享答案。

就我而言,解决方案非常简单,我已经在anaconda env中安装了mmdnn。并且文件{'A': {1: 2, 2: 3, 3: 4}, 'B': {2: 44.0}} 存储在以下位置:

“用户名/ anaconda3 / envs / env_name / bin / mmtoir”

此文件中需要明确引用我的神经网络类。在您的情况下,您需要:

1)找到该文件

2)复制CIFAR类的定义

mmtoir

放入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 文件中。保存并再次运行命令。

以上是解决此问题的两种方法之一,相对简单。您可以阅读此discussion了解更多详细信息。

注意:MMdnn当前仅支持mmtoir。因此,请确保您的模型也使用相同的版本进行训练。使用PyTorch=0.4.0,所有步骤都可以正常运行,但最终会引发错误。