Currently, we have already implemented both the the PyTorch -> IR part and the IR -> PyTorch part.
√ - Correctness tested
o - Some difference after conversion
space - not tested
The PyTorch parser is modified from branch pytorch , using jit CppOP to build the graph.
Any contribution is welcome.
Extract PyTorch pre-trained models
You can refer PyTorch model extractor to extract your pytorch models.
$ mmdownload -f pytorch -h Supported models: ['alexnet', 'densenet121', 'densenet161', 'densenet169', 'densenet201', 'inception_v3', 'resnet101', 'resnet152', 'resnet18', 'resnet34', 'resnet50', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', 'vgg19', 'vgg19_bn'] $ mmdownload -f pytorch -n resnet101 -o ./ Downloading: "https://download.pytorch.org/models/resnet101-5d3b4d8f.pth" to /my/home/.torch/models/resnet101-5d3b4d8f.pth ███████████████████| 102502400/102502400 [00:06<00:00, 15858546.50it/s] PyTorch pretrained model is saved as [./imagenet_resnet101.pth].
Convert Pytorch pre-trained models to IR
You can convert the whole pytorch model to IR structure. Please remember for the generality, we now only take the whole model
pth, not just the state dict. To be more specific, it is save using
torch.load() can load the whole model.
$ mmtoir -f pytorch -d resnet101 --inputShape 3,224,224 -n imagenet_resnet101.pth
Please bear in mind that always add
--inputShape argparse. This thing is different from other framework because pytorch is a dynamic framework.
Then you will get
IR network structure is saved as [resnet101.json]. IR network structure is saved as [resnet101.pb]. IR weights are saved as [resnet101.npy].
Convert models from IR to PyTorch code snippet and weights
You can use following bash command to convert the IR architecture file [inception_v3.pb] and weights file [inception_v3.npy] to Caffe Python code file[pytorch_inception_v3.py] and IR weights file suit for caffe model[pytorch_inception_v3.npy]
Note: We need to transform the IR weights to PyTorch suitable weights. Use argument -dw to specify the output weight file name.
$ mmtocode -f pytorch -n inception_v3.pb --IRWeightPath inception_v3.npy --dstModelPath pytorch_inception_v3.py -dw pytorch_inception_v3.npy Parse file [inception_v3.pb] with binary format successfully. Target network code snippet is saved as [pytorch_inception_v3.py]. Target weights are saved as [pytorch_inception_v3.npy].
Generate PyTorch model from code snippet file and weight file
You can use following bash command to generate PyTorch model file [pytorch_inception_v3.pth] from python code [pytorch_inception_v3.py] and weights file [pytorch_inception_v3.npy] for further usage.
$ mmtomodel -f pytorch -in pytorch_inception_v3.py -iw pytorch_inception_v3.npy -o pytorch_inception_v3.pth PyTorch model file is saved as [pytorch_inception_v3.pth], generated by [pytorch_inception_v3.py] and [pytorch_inception_v3.npy]. Notice that you may need [pytorch_inception_v3.py] to load the model back.
Detail scripts of Tensorflow slim resnet_v1_101 model to PyTorch conversion are in issue 22. You can refer it to implement your conversion.
Ubuntu 16.04 with
- PyTorch 0.4.0
The main dataflow in a pytorch network is converted from NHWC(channel last) to NCHW(channel first) format, but some operators (like Concat) with axis may not transform correctly. You may need to correct it manually.
Currently, no RNN-related operations supported
- There are two types models saved in PyTorch. One is including architecture and weights, which is supported in the MMdnn now. The other one is only including the weights, which is not supported now.
only_weight_file = "./alexnet-owt-4df8aa71.pth" # Download from the model zoo architecture_weight_file = "imagenet_alexnet.pth" # Download using mmdownload() m = torch.load(only_weight_file) # <class 'collections.OrderedDict'> m_1 = torch.load(architecture_weight_file) # <class 'torchvision.models.alexnet.AlexNet'> supported!
- When you get the error "AttributeError: 'collections.OrderedDict' object has no attribute 'state_dict'" , it's because you use the model only include weights part. You need to save a new model with archietecture
- How to load the converted PyTorch model ?
import torch import imp import numpy as np MainModel = imp.load_source('MainModel', "tf_pytorch_vgg19.py") the_model = torch.load("tf_pytorch_vgg19.pth") the_model.eval() x = np.random.random([224,224,3]) x = np.transpose(x, (2, 0, 1)) x = np.expand_dims(x, 0).copy() data = torch.from_numpy(x) data = torch.autograd.Variable(data, requires_grad = False).float() predict = the_model(data)