pytorch转caffe步骤

韩耘豪

2023-12-01

写在前面

来记录一下把用pytorch训练好的模型转成caffe去预测的步骤框架，代码只展现主要部分~

步骤

保存pytorch参数名和权重。先把pytorch的参数名和权重存成词典，例如存到npy文件里–>.得到npy文件；
合并卷积层和bn层的参数（非必须）。将卷积层和batch normalization层的权重进行融合，更新npy文件里卷积层的权重–>.得到npy文件；
建立caffe的.prototxt文件。即按照pytorch的网络架构写出caffe的架构–>得到.prototxt文件；
建立参数名映射。建立caffe—>pytorch的参数名称映射，可以存在字典或者txt里面–>得到.txt文件；
建立caffemodel。按照参数名称的映射，将pytorch的参数权重赋予对应的caffe参数–>得到.caffemodel文件；
推测。在python中将caffemodel嵌进去，进行推断，看看能不能得到对应结果。

1. 保存pytorch参数名和权重（.npy）

用torch将模型和参数加载好之后，用model.state_dict().items()就可以把模型model中的参数名和权重拎出来了。其实这一步用model.named_parameters()也可以，但是这种方法没有把bn层的running_mean和running_var，而在融合卷积层和bn层中要用到，所以我用了前一个方法。

import torch
class Model(nn.Module):
'''定义一个Model()类的模型，pytorch定义模型常规操作'''
    def __init__(self):
    '''定义层'''
        self.deconv1 = nn.ConvTranspose2d(...)
        ...
    def forward(self, x):
    '''定义实际的层操作'''
        score = self.deconv1(x)
        ...
        return score
model_file = "xxx.ckpt" # 训练好的模型参数
model = Model(..) # 建立Model类的模型
model.load_state_dict(torch.load(model_file, map_location='cpu'))
model.eval() #将模型设置为验证模式，这会将dropout关掉，将bn固定

weights = dict() # 储存模型的参数名和权重
for k,v in model.state_dict().items():
    weights[k] = v
np.save("params.npy", weights)

2. 合并卷积层和bn层的参数（.npy）

这一步非必须，适用于卷积层后紧跟bn层的情形，主要是为了节省计算量，可参考github:Merge_BN中的操作，实际上bn层的weight和bias就是bn层的gamma和beta，通过公式展开就可以得到新的weight和bias了。

3. 建立prototxt文件（.prototxt）

个人觉得主要是两种写法。

一种是从0建立prototxt，这个fully convolution network的原装代码就是这种方法fcn-berkeleyvision。先用n = caffe.NetSpec()建立网络n（这个n为caffe.net_spec.NetSpec类），然后对n.xxx赋予网络操作如L.Convolution(...)定义层名xxx的层。

一种是在之前网络的基础上进行增加或修改，根据stackoverflow:如何用pycaffe重建层，我定义了一个函数add_layer：

import caffe
from caffe import layers as L, params as P
from google.protobuf import text_format
import caffe.proto.caffe_pb2 as caffe_pb2

def add_layer(net, layer, name, bottom, top=None):
    """
    net: the base net. type: caffe_pb2.NetParameter()
    layer: caffe.layers. type: caffe.net_spec.Top
    name: layer name. type: str
    bottom: bottom layer. type: list containing all the bottom names.
    top: top layer. type: list containing all the top names.
    """
    l = net.layer.add()
    l.CopyFrom(layer.to_proto().layer[0])
    l.name = name
    for b in bottom: # l.CopyFrom之后bottom默认就是空的，直接用append往里面加就好了
        l.bottom.append(b)
    if top:
        l.top[:] = [] # l.CopyFrom之后会初始化top，先把它们清空
        for t in top:
            l.top.append(t)
    else: # 不设置top的时候，和name一样
        l.top[0] = name

# 例如vgg16是基础网络
# vgg16.prototxt下载：https://github.com/soeaver/caffe-model/blob/master/cls/vgg/deploy_vgg16-pytorch.prototxt
deploy_path = "vgg16.prototxt"
n = caffe_pb2.NetParameter() # 这个n属于caffe_pb2.NetParameter类
text_format.Merge(open(deploy_path).read(), n)

# 瞎写几层
add_layer(net=n, layer=L.Deconvolution(convolution_param=dict(num_output=512, kernel_size=3, stride=2, pad=0, dilation=1, bias_term=True)), name="g_deconv1", bottom=["pool5"])
...
add_layer(net=n, layer=L.InnerProduct(num_output=2), name="p_output", bottom=["p_fc2"])

with open("newmodel.prototxt","w") as f:
    f.write(str(n))

这样就可以在vgg16.prototxt上加层，储存到newmodel.prototxt中啦。

实际上pycaffe定义的这些layer（例如L.Deconvolution）都是caffe.net_spec.NetSpec类，而NetSpec.to_proto则得到NetParameter类，所以要么就从一开始直接建立NetSpec类，然后直接加layer，要么就一开始建立NetParameter类，然后往里面加layer.to_proto。

4. 建立caffe—>pytorch的参数名称映射（.txt）

接下来把caffe里面需要赋权重的层名，和对应Pytorch里面的层名对应起来就好了，比如我是建立了一个txt文件，空格左边是pytorch中的层名，右边是caffe中我起的层名。形式也很随意，只要能一一对应上就好啦，比如用词典dict去存也是可以的。

pretrained_net.deconv5 g_deconv5
pretrained_net.deconv2 g_deconv2
conv0_2 p_conv0_2
...

5. 按照参数名称映射赋予权重，得到caffemodel

有了caffe的网络架构（.prototxt）、pytorch的权重（.npy）和pytorch与caffe的层名对应关系（.txt），就可以得到caffemodel了！（激动地搓手手）把上面的参数名称映射读成词典，然后对应地把pytorch里面层的权重赋给caffe的层就好了。

deploy_file = "xxx.prototxt" # caffe的网络架构
weights_file = "xxx.npy" # pytorch的权重
mapping_file = "xxx.txt" # pytorch与caffe的层名映射
caffemodel_file = "xxx.caffemodel" # 这个是即将要储存的caffemodel名

def load_weights(net, weights, mapping):
    for key, value in net.params.items():
        if key in mapping.keys():
            key_weights = mapping[key] + '.weight'
            key_biases = mapping[key] + '.bias'
            net.params[key][0].data[...] = weights[key_weights].cpu().detach()
            net.params[key][1].data[...] = weights[key_biases].cpu().detach()
            print("Successfully load {}-->{}.".format(mapping[key], key))
        else:
            print('Not mapped: ', key)
    return net

mapping = {}
mapping = ...# 这一步略掉，就是将pytorch的层名存为mapping的key，caffe的层名存为val
net = caffe.Net(deploy_file, caffe.TEST)
weights = np.load(weights_file).tolist()
net = load_weights(net, weights, mapping)
net.save(caffemodel_file)

6. 用转换后的caffemodel预测

转换好caffemodel之后，就可以在python中用它测试一下啦！

import caffe

# 通过以下方法读入网络结构和参数
deploy_file = "xxx.prototxt"
caffemodel_file = "xxx.caffemodel"
net = caffe.Net(deploy_file, caffemodel_file, caffe.TEST)

# 读入数据
im = ...# 此处为数据处理过程，假设读入图片，经处理之后得到im
net.blobs['data'].data[...] = im 

# 跑起来吧，caffe！
net.forward()

# 把数据经过xxx层后的结果输出来
out = net.blobs['xxx'].data[0]

后记：pytorch到caffe的坑

目前我遇到的问题除了自己坑自己，其他就主要是发生在层定义上，也就是写prototxt的时候发现caffe的一些定义和pytorch会有所不同。

反卷积的padding

pytorch的反卷积是有output_padding的设置的，也就是在反卷积的时候，会先对输入进行一次padding，反卷积之后，对输出再进行一次单边的padding，但caffe只有输入padding的设置（也就是caffe反卷积api里面的pad）。也有人在pytorch转caffe的github上问过类似问题：pytorch2caffe:issues7。原本我想的很简单，既然caffe没有这个设置，那对于输出做一下padding就好了嘛！比如有个padding层之类的，但搜了一圈没有搜着比较简单的。

我目前想到的解决方法只有这两个，要是有更好方法的小可爱欢迎提出！

github上有一些padding层的方法，如caffe原装代码里面的issue6506。但是要再配置caffe。
对网络进行一些微调，让它不被这种问题困扰，例如可以加个crop层之类的，比如我遇到这个问题是因为我对某一层进行反卷积之后，要和另一层的图相加融合，如果反卷积之后的图和另一层的图尺寸不对齐，就会发生报错，所以可以进行一些微调，比如把padding置为0，让它反卷积之后的图比另一层的图稍大一点，然后对反卷积后的图用Crop层进行裁剪，使其大小与另一层的图一致。这样做相当于改了网络结构，所以最好要么重新按这个架构再训练一次pytorch之后再进行权重的迁移，要么就直接改成训练caffe了。

maxpooling的尺寸大小

maxpooling的时候发生了caffe比pytorch得到的尺寸大一点的情形，主要还是取整的时候caffe往上取而pytorch往下取，pytorch的maxpool可以设置ceil_mode，把它设置为True就好了。但是因为改的是pytorch结构，和之前的参数权重尺寸不同了，所以要重新训练pytorch模型。