pytorch转caffe步骤

写在前面

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

步骤

1. 保存pytorch参数名和权重。先把pytorch的参数名和权重存成词典，例如存到npy文件里–>.得到npy文件；
2. 合并卷积层和bn层的参数（非必须）。将卷积层和batch normalization层的权重进行融合，更新npy文件里卷积层的权重–>.得到npy文件；
3. 建立caffe的.prototxt文件。即按照pytorch的网络架构写出caffe的架构–>得到.prototxt文件；
4. 建立参数名映射。建立caffe—>pytorch的参数名称映射，可以存在字典或者txt里面–>得到.txt文件；
5. 建立caffemodel。按照参数名称的映射，将pytorch的参数权重赋予对应的caffe参数–>得到.caffemodel文件；
6. 推测。在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