PyTorch转Caffe模型

最新推荐文章于 2024-07-09 11:35:54 发布

xiaohe9275

最新推荐文章于 2024-07-09 11:35:54 发布

阅读量2.2k

点赞数 2

分类专栏： PyTorch Caffe 文章标签： caffe pytorch 深度学习

本文链接：https://blog.csdn.net/xiaohe9275/article/details/121981794

版权

PyTorch 同时被 2 个专栏收录

3 篇文章 0 订阅

订阅专栏

Caffe

3 篇文章 0 订阅

订阅专栏

PyTorch转Caffe模型

使用前需要将caffe.proto和caffe_pb2.py替换成自己的。

一、代码执行入口：xxx_pytorch_to_caffe.py

以alexnet_pytorch_to_caffe.py为例，代码如下：

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

if __name__=='__main__':
    name='alexnet'
    net=alexnet(True)
    input=Variable(torch.ones([1,3,226,226]))
    pytorch_to_caffe.trans_net(net,input,name)
    pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name))
    pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name))

这块是整个转模型代码的入口，首先定义自己的网络对象和一个输入（假数据），调用trans_net生成xxx.prototxt和xxx.caffemodel，最后保存。

二、主体功能代码：pytorch_to_caffe.py

1、trans_net

def trans_net(net,input_var,name='TransferedPytorchModel'):
    print('Starting Transform, This will take a while')
    log.init([input_var])
    log.cnet.net.name=name
    log.cnet.net.input.extend([log.blobs(input_var)])
    log.cnet.net.input_dim.extend(input_var.size())
    global NET_INITTED
    NET_INITTED=True
    for name,layer in net.named_modules():
        layer_names[layer]=name
    print("torch ops name:", layer_names)
    out = net.forward(input_var)
    print('Transform Completed')

log对象记录了转换的caffe模型中各层的名字、网络结构、模型参数等信息，起初先填写xxx.prototxt的输入dim信息。
log的类如下：

class TransLog(object):
    def __init__(self):
        """
        doing init() with inputs Variable before using it
        """
        self.layers={}
        self.detail_layers={}  
        self.detail_blobs={}  
        self._blobs=Blob_LOG()
        self._blobs_data=[]
        self.cnet=caffe_net.Caffemodel('')
        self.debug=True

    def init(self,inputs):
        """
        :param inputs: is a list of input variables
        """
        self.add_blobs(inputs)
    #后面代码省略。。。

然后调用net对象的forward函数执行一次推理，在推理过程中会按照网络结构分别调用各层的底层函数，如运算到卷积层（nn.Conv2d）时，会调用torch.nn.functional.conv2d，但在pytorch_to_caffe.py中torch.nn.functional.conv2d被替换成了自定义的函数_conv2d，在该函数中会根据pytorch中定义的卷积层信息生成caffe的卷积层网络结构信息，同时保存卷积层权重和偏置参数：

def _conv2d(raw,input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
    #通过卷积层的输入blob对应的标识符查找其blob名并打印，该blob名会在上一层中被记录
    print('conv: ',log.blobs(input))
    x=raw(input,weight,bias,stride,padding,dilation,groups)
    name=log.add_layer(name='conv')
    #记录该卷积层输出blob名
    log.add_blobs([x],name='conv_blob')
    #定义caffe的xxx.prototxt中的网络层
    layer=caffe_net.Layer_param(name=name, type='Convolution',
                                bottom=[log.blobs(input)], top=[log.blobs(x)])
    layer.conv_param(x.size()[1],weight.size()[2:],stride=_pair(stride),
                     pad=_pair(padding),dilation=_pair(dilation),bias_term=bias is not None,groups=groups)
    #保存卷积层权重和偏置参数
    if bias is not None:
        layer.add_data(weight.cpu().data.numpy(),bias.cpu().data.numpy())
    else:
        layer.param.convolution_param.bias_term=False
        layer.add_data(weight.cpu().data.numpy())
    log.cnet.add_layer(layer)
    return x

F.conv2d=Rp(F.conv2d,_conv2d)

2、Rp

# 核心组件，通过该类，实现对torch的function中的operators的输入，输出以及参数的读取
class Rp(object):
    def __init__(self,raw,replace,**kwargs):
        # replace the raw function to replace function
        self.obj=replace
        self.raw=raw

    def __call__(self,*args,**kwargs):
        if not NET_INITTED:
            return self.raw(*args,**kwargs)
        for stack in traceback.walk_stack(None):
            if 'self' in stack[0].f_locals:
                layer=stack[0].f_locals['self']
                if layer in layer_names:
                    log.pytorch_layer_name=layer_names[layer]
                    print(layer_names[layer])
                    break
        out=self.obj(self.raw,*args,**kwargs)
        # if isinstance(out,Variable):
        #     out=[out]
        return out

实现自定义函数和torch.nn.functional中的函数的替换，如：

F.conv2d=Rp(F.conv2d,_conv2d)
F.linear=Rp(F.linear,_linear)
F.relu=Rp(F.relu,_relu)
F.leaky_relu=Rp(F.leaky_relu,_leaky_relu)
F.max_pool2d=Rp(F.max_pool2d,_max_pool2d)
F.avg_pool2d=Rp(F.avg_pool2d,_avg_pool2d)
F.adaptive_avg_pool2d = Rp(F.adaptive_avg_pool2d,_adaptive_avg_pool2d)
F.dropout=Rp(F.dropout,_dropout)
F.threshold=Rp(F.threshold,_threshold)
F.prelu=Rp(F.prelu,_prelu)
F.batch_norm=Rp(F.batch_norm,_batch_norm)
F.instance_norm=Rp(F.instance_norm,_instance_norm)
F.softmax=Rp(F.softmax,_softmax)
F.conv_transpose2d=Rp(F.conv_transpose2d,_conv_transpose2d)
F.interpolate = Rp(F.interpolate,_interpolate)
F.sigmoid = Rp(F.sigmoid,_sigmoid)
F.tanh = Rp(F.tanh,_tanh)
F.tanh = Rp(F.tanh,_tanh)
F.hardtanh = Rp(F.hardtanh,_hardtanh)

此外，有些版本的pytorch中的个别层底层调用的不是torch.nn.functional中的函数，而是torch，如sigmoid等，所以此时就要替换torch中的函数，如：

torch.split=Rp(torch.split,_split)
torch.max=Rp(torch.max,_max)
torch.cat=Rp(torch.cat,_cat)
torch.div=Rp(torch.div,_div)

三、遇到过的问题

1、ModuleNotFoundError: No module named 'google'
原因：python中没有安装protobuf

2、AttributeError: 'PoolingParameter' object has no attribute 'ceil_mode'
原因：使用的caffe中pooling层没有ceil_mode参数
解决方法：在caffe的pooling层添加该参数及相应的源码
1）在caffe.protode PoolingParameter中添加ceil_mode
2）修改pooling_layer.hpp中PoolingLayer类

3、TypeError: _avg_pool2d() takes from 3 to 7 postinal arguments but 8 were given
原因：avg pool层参数不对
解决方法：修改pytorch_to_caffe.py中_avg_pool2d

#def _avg_pool2d(raw,input, kernel_size, stride = None, padding = 0, ceil_mode = False, count_include_pad = True):
def _avg_pool2d(raw,input, kernel_size, stride = None, padding = 0, ceil_mode = False, count_include_pad = True, divisor_override = None):
    x = raw(input, kernel_size, stride, padding, ceil_mode, count_include_pad)
    _pool('ave',raw,input, x, kernel_size, stride, padding,ceil_mode)
    return x

4、如果训练时模型开了多卡训练，推理时也需要这样，即net = nn.DataParallel(net).cuda()，但由于数据被拷贝了多份，第一个卷积层拿到的blob地址不再是输入的了，导致找不到这个blob。
解决方法：去掉net = nn.DataParallel(net).cuda()，且将
net.load_state_dict(checkpoint)
checkpoint = torch.load("xxx.ckpt")
换成net.load_state_dict({k.replace('module.',''):v for k,v in torch.load("xxx.ckpt").items()})

5、当前版本的sigmoid是通过torch.sigmoid()实现的，不是F.sigmoid，所以需要添加torch.sigmoid = Rp(torch.sigmoid,_sigmoid)，原先只有F.sigmoid = Rp(F.sigmoid,_sigmoid)

6、layer_name = log.add_layer(name="expand_as")#, with_num=True)，log.add_layer参数没有with_num

7、当前的caffe不支持双线性插值层，自定义了一个插值层。重新定义_interpolate2()来取代F.interpolate，即：F.interpolate = Rp(F.interpolate,_interpolate2)，_interpolate2()代码如下：

def _interpolate2(raw, input,size=None, scale_factor=None, mode='nearest', align_corners=None):
    x = raw(input,size , scale_factor ,mode)

    layer_name = log.add_layer(name='interpolate')
    top_blobs = log.add_blobs([x], name='interpolate_blob'.format(type))
    layer = caffe_net.Layer_param(name=layer_name, type='Interp',
                                  bottom=[log.blobs(input)], top=top_blobs)

    layer.interp_param(scale_factor= scale_factor)
    log.cnet.add_layer(layer)
    return x

8、1.9版本pytorch的torch.cat中的维度参数名称为dim，不是dimension，因此将def _cat(raw,inputs, dimension=0)换成def _cat(raw,inputs, dim=0)。