Pytorch轉Caffe

https://blog.csdn.net/vvnzhang2095/article/details/91439924

 

pytorch 與caffe2的互用

https://blog.csdn.net/weixin_40671425/article/details/90700908

寫在前面

來記錄一下把用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模型。