hahnyuan / nn_tools

Neural Network Tools: Converter and Analyzer. For caffe, pytorch, draknet and so on.
MIT License
355 stars 63 forks source link

Incomplete exported ResNet18 prototxt from PyTorch #15

Closed xpngzhng closed 5 years ago

xpngzhng commented 5 years ago

PyTorch version 1.1.0 torchvision version 0.2.0

I add the following two lines after https://github.com/hahnyuan/nn_tools/blob/master/example/testify_pytorch_to_caffe_example.py#L60

if name not in blob2layer:
    continue

then I can get perfect alexnet model conversion from pytorch to caffe.

But I failed in converting resnet18. I get the following error message:

F0627 11:57:11.057982 50772 insert_splits.cpp:35] Unknown bottom blob 'add_blob10' (layer 'layer1_0_relu_sub1', bottom index 0)

then I checked the resnet18 prototxt, add_blob10 only apprears once:

layer {
  name: "layer1_0_bn2"
  type: "BatchNorm"
  bottom: "conv_blob8"
  top: "batch_norm_blob9"
  batch_norm_param {
    use_global_stats: true
    eps: 1e-05
  }
}
layer {
  name: "noname_bn_scale11"
  type: "Scale"
  bottom: "batch_norm_blob9"
  top: "batch_norm_blob9"
  scale_param {
    bias_term: true
  }
}
layer {
  name: "layer1_0_relu_sub1"
  type: "ReLU"
  bottom: "add_blob10"
  top: "relu_blob11"
}

the full resnet18 prototxt is here resnet18.txt All the element wise add operations are missing

draw the net in netscope, I see fragmented nets http://ethereon.github.io/netscope/#/editor 图片

For your information, the resnet_pytorch_to_caffe.py and testify_pytorch_to_caffe_example.py are as follows, somewhat different from the original versions. I did not load the pretrained model, but save the pth model in the current directory for later testing.

import sys
sys.path.insert(0,'.')
import torch
from torch.autograd import Variable
from torchvision.models.resnet import *
import pytorch_to_caffe

if __name__=='__main__':
    name='resnet18'
    net=resnet18(False)
    net.eval()
    torch.save(net.state_dict(), 'resnet18.pth')
    input=Variable(torch.ones([1,3,224,224]))
    pytorch_to_caffe.trans_net(net,input,name)
    pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name))
    pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name))
import sys
import caffe
import torch
import numpy as np
import argparse
from collections import OrderedDict
from torch.autograd import Variable
import torch.nn as nn
import pytorch_to_caffe

def arg_parse():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', '-m', default='alexnet')
    parser.add_argument('--decimal', '-d', default=2)
    parser.add_argument('--gpu', '-gpu', action='store_true')
    args = parser.parse_args()
    return args

def generate_random(shape, gpu=False):
    data_np = np.random.rand(np.prod(shape)).reshape(shape)
    data_torch = Variable(torch.Tensor(data_np))
    if gpu:
        data_torch = data_torch.cuda()
    return [data_np], [data_torch]

def get_input_size(caffe_net):
    input_name = caffe_net.inputs[0]
    return caffe_net.blobs[input_name].data.shape

def forward_torch(net, data):
    blobs = OrderedDict()
    module2name = {}
    for layer_name, m in net.named_modules():
        layer_name = layer_name.replace('.', '_')
        module2name[m] = layer_name
        # turn off all the inplace operation
        if hasattr(m, 'inplace'):
            m.inplace = False

    def forward_hook(m, i, o):
        o_np = o.data.cpu().numpy()
        blobs[module2name[m]] = o_np

    for m in net.modules():
        m.register_forward_hook(forward_hook)
    output = net.forward(*data)
    if isinstance(output, tuple):
        outputs = []
        for o in output:
            outputs.append(o.data.cpu().numpy())
    else:
        outputs = [output.data.cpu().numpy()]
    return blobs, outputs

def forward_caffe(net, data):
    for input_name, d in zip(net.inputs, data):
        net.blobs[input_name].data[...] = d
    rst = net.forward()
    blobs = OrderedDict()
    blob2layer = {}
    print(type(net.top_names))
    for layer_name, tops in net.top_names.items():
        for top in tops:
            blob2layer[top] = layer_name
    for k, v in blob2layer.items():
        print(k, v)
    for name, value in net.blobs.items():
        print('----', name)
        if name not in blob2layer:
            continue
        layer_name = blob2layer[name]
        value = value.data
        if layer_name in blobs:
            blobs[layer_name].append(value)
        else:
            blobs[layer_name] = [value]
    outputs = []
    for output_name in net.outputs:
        outputs.append(rst[output_name])
    return blobs, outputs

def test(net_caffe, net_torch, data_np, data_torch, args):
    blobs_caffe, rsts_caffe = forward_caffe(net_caffe, data_np)
    blobs_torch, rsts_torchs = forward_torch(net_torch, data_torch)
    # test the output of every layer
    for layer, value in blobs_caffe.items():
        if layer in blobs_torch:
            value_torch = blobs_torch[layer]
            value = value[0]
            if value.size != value_torch.size: continue
            if 'relu' in layer: continue
            try:
                np.testing.assert_almost_equal(value, value_torch, decimal=args.decimal)
                print("TEST layer {}: PASS".format(layer))
            except:
                print("TEST layer {}: FAIL".format(layer))
                # np.testing.assert_almost_equal(np.clip(value, min=0), np.clip(value_torch, min=0))
    # test the output
    print("TEST output")
    for rst_caffe, rst_torch in zip(rsts_caffe, rsts_torchs):
        np.testing.assert_almost_equal(rst_caffe, rst_torch, decimal=args.decimal)
    print("TEST output: PASS")

if __name__ == '__main__':
    args = arg_parse()

    if args.model == 'alexnet':
        # Alexnet example
        from torchvision.models.alexnet import alexnet

        net_torch = alexnet(False).eval()
        net_torch.load_state_dict(torch.load('alexnet.pth'))
        if args.gpu:
            net_torch.cuda()
        try:
            net_caffe = caffe.Net('alexnet.prototxt', 'alexnet.caffemodel', caffe.TEST)
        except:
            raise ("Please run alexnet_pytorch_to_caffe.py first")
        shape = get_input_size(net_caffe)
        data_np, data_torch = generate_random(shape, args.gpu)
        test(net_caffe, net_torch, data_np, data_torch, args)
    elif args.model == 'resnet18':
        # ResNet example
        from torchvision.models.resnet import resnet18

        net_torch = resnet18(False).eval()
        net_torch.load_state_dict(torch.load('resnet18.pth'))
        if args.gpu:
            net_torch.cuda()
        net_caffe = caffe.Net('resnet18.prototxt', 'resnet18.caffemodel', caffe.TEST)
        shape = get_input_size(net_caffe)
        data_np, data_torch = generate_random(shape, args.gpu)
        test(net_caffe, net_torch, data_np, data_torch, args)
    elif args.model == 'inception_v3':
        # Inception_v3 example
        from torchvision.models.inception import inception_v3

        net_torch = inception_v3(False, transform_input=False).eval()
        if args.gpu:
            net_torch.cuda()
        net_caffe = caffe.Net('inception_v3.prototxt', 'inception_v3.caffemodel', caffe.TEST)
        shape = get_input_size(net_caffe)
        data_np, data_torch = generate_random(shape, args.gpu)
        test(net_caffe, net_torch, data_np, data_torch, args)
    elif args.model == 'densenet121':
        # Densnet example
        from torchvision.models.densenet import densenet121

        net_torch = densenet121(False).eval()
        if args.gpu:
            net_torch.cuda()
        net_caffe = caffe.Net('densenet121.prototxt', 'densenet121.caffemodel', caffe.TEST)
        shape = get_input_size(net_caffe)
        data_np, data_torch = generate_random(shape, args.gpu)
        test(net_caffe, net_torch, data_np, data_torch, args)
    else:
        raise NotImplementedError()
buleZia commented 5 years ago

Did you slove this problem? I met this problem when i used pytorch0.4.1 & torchvision 0.2.1 converting resnet18

xpngzhng commented 5 years ago

@buleZia It may be a mismatch between the caffe.proto of the caffe installed in your computer and that in Caffe directory inside this repo. Follow the advice in Caffe/ReadMe.md, replace the caffe.proto with yours and recompile to get caffe_pb2.py