調試從Caffe模型轉換過來的Pytorch模型

系列博客目錄：Caffe轉Pytorch模型系列教程 概述

相信看完本系列的前兩博客，如果你運氣好，你的Pytorch模型說不定已經可以工作了，那麼Congratulations，不必往下看了。如果結果不對，那麼，你就需要看看我慘痛的調試經歷了。

• 本文用的Caffe網絡模型文件：SfSNet_deploy.prototxt（右鍵另存爲）。
• SfSNet的Pytorch代碼地址：model.py(右鍵另存爲，來源於上一篇博客)。
• SfSNet的權重：SfSNet.caffemodel.h5（右鍵另存爲）。

一、檢查Pytorch模型是否正確地載入了參數

這個就很簡單了，最簡單的思路就是，把網絡的參數打印出來，再和從.caffemodel提取出來的對比（這是最容易想到的辦法，我一開始也是用的這種方法）。

# coding=utf-8
from __future__ import absolute_import, division, print_function
import pickle
from src.models.model import SfSNet

if __name__ == '__main__':
    # 載入提取到的權重
    f = open('SfSNet-Caffe/weights.pkl', 'rb')
    name_weights = pickle.load(f)
    f.close()
    # 新建網絡實例
    net = SfSNet()
    # 載入參數
    net.load_weights_from_pkl('SfSNet-Caffe/weights.pkl')
    # 設置爲測試模式
    net.eval()

    params = dict(net.named_parameters())

    # 檢查conv1的參數是否載入正確
    print(params['conv1.weight'].detach().numpy())
    print(name_weights['conv1']['weight'])

但是這種方法，對眼睛真的超超超不友好啊，所以稍微改進下，寫函數替我們去判斷：

import numpy as np
def same(arr1, arr2):
    # type: (np.ndarray, np.ndarray) -> bool
    # 判斷shape是否相同
    assert arr1.shape == arr2.shape
    # 對應元素相減求絕對值
    diff = np.abs(arr1 - arr2)
    # 判斷是否有任意一個兩元素之差大於閾值1e-5
    return (diff < 1e-5).any()


def compare(layer, params, name_weights):
    # type: (str, dict, dict) -> tuple[bool, bool]
    # 判斷權重
    w = same(params[layer+'.weight'].detach().numpy(), name_weights[layer]['weight'])
    # 判斷偏置
    b = same(params[layer+'.bias'].detach().numpy(), name_weights[layer]['bias'])
    return w, b

然後就可以使用：

print(compare('conv1', params, name_weights))

來檢查conv1的參數。其他層的參數檢查我就不一一舉例了，讀者知道怎麼辦的。

二、檢查網絡模型是否正確

1、對比Pytorch模型與.prototxt可視化結果

就以本文的例子來說，就是拿着SfSNet_deploy.prototxt的可視化結果（可視化方法參見：把Caffe的模型轉換爲Pytorch模型）和Pytorch模型的對比。如果幸運，檢查到哪裏出錯了，修改完，模型ok了，那完美；如果還是不行，看下一步。

2、運行兩個網絡對應的層並對比結果

這一步的目的是檢查問題到底出在那一層。知道出在那一層之後，纔好解決bug。

2.1 Caffe模型獲取指定層的結果

# coding=utf-8
from __future__ import absolute_import, division, print_function
import caffe

if __name__ == '__main__':
    # prototxt文件
    MODEL_FILE = 'SfSNet-Caffe/SfSNet_deploy.prototxt'
    # 預先訓練好的caffe模型
    PRETRAIN_FILE = 'SfSNet-Caffe/SfSNet.caffemodel.h5'
    # 載入網絡
    net = caffe.Net(MODEL_FILE, PRETRAIN_FILE, caffe.TEST)
    ... 
    # 前向傳播
    out = net.forward(end='conv3')

forward函數有個end參數，設置了此參數時，網絡在向前傳播到conv3時（遇到第一個conv3就會立即返回結果），就會返回conv3的結果：一個dict，key是’conv3’，值是conv3的結果。

2.2 Pytorch模型獲取指定層結果

這個很簡單了，直接修改model.py，在forward函數裏返回該層的結果就行。

class SfSNet(nn.Module):  # SfSNet = PS-Net in SfSNet_deploy.prototxt
    def __init__(self):
        # C64
        super(SfSNet, self).__init__()
        # TODO 初始化器 xavier
        self.conv1 = nn.Conv2d(3, 64, 7, 1, 3)
        self.bn1 = nn.BatchNorm2d(64)
        # C128
        self.conv2 = nn.Conv2d(64, 128, 3, 1, 1)
        self.bn2 = nn.BatchNorm2d(128)
        # C128 S2
        self.conv3 = nn.Conv2d(128, 128, 3, 2, 1)
        ...

    def forward(self, inputs):
        # C64
        x = F.relu(self.bn1(self.conv1(inputs)))
        # C128
        x = F.relu(self.bn2(self.conv2(x)))
        # C128 S2
        conv3 = self.conv3(x)
        # 返回conv3的結果
        return conv3
        # ------------RESNET for normals------------
        # RES1
        x = self.n_res1(conv3)
        ...
        return normal, albedo, light

    def load_weights_from_pkl(self, weights_pkl):
        from torch import from_numpy
        with open(weights_pkl, 'rb') as wp:
            ...
            state_dict = {}
            ...
            self.load_state_dict(state_dict)

然後：

if __name__ == '__main__':
    # 新建網絡實例
    net = SfSNet()
    # 載入參數
    net.load_weights_from_pkl('SfSNet-Caffe/weights.pkl')
    # 設置爲測試模式
    net.eval()
	# 預處理圖像
    image = ...
    # 前向傳播
    out = net(image)

2.3 準備相同的預處理代碼

這一步相當重要，數據輸入都不一樣，你咋讓兩個網絡輸出同樣的結果呢？？？定義SfSNet的預處理函數：

import cv2
def read_image(path):
    # type: (str) -> np.ndarray
    # 讀取圖像
    image = cv2.imread(path)
    # 調整大小爲SfSNet的代碼
    image = cv2.resize(image, (128, 128))
    # 縮放到0~1
    image = np.float32(image)/255.0
    # (128, 128, 3) to (3, 128, 128)
    image = np.transpose(image, [2, 0, 1])
    # (128, 128, 3) to (1, 3, 128, 128)
    image = np.expand_dims(image, 0)
    
    return image

不同的模型可能有不同的預處理步驟，這只是SfSNet的預處理步驟。

2.4 獲取conv3的結果

• Caffe模型代碼：

if __name__ == '__main__':

    # prototxt文件
    MODEL_FILE = 'SfSNet-Caffe/SfSNet_deploy.prototxt'
    # 預先訓練好的caffe模型
    PRETRAIN_FILE = 'SfSNet-Caffe/SfSNet.caffemodel.h5'
    # 定義網絡
    net = caffe.Net(MODEL_FILE, PRETRAIN_FILE, caffe.TEST)
    # 讀取並預處理圖像
    im = read_image('data/1.png_face.png')
    # 前向傳播
    out = net.forward(end='conv3')

    print(out.keys())
    # 保存
    np.save('conv3.caffe.npy', out['conv3'])

• Pytorch模型代碼：
  先修改model.py的forward函數，然後再運行代碼：

if __name__ == '__main__':
    # 新建網絡實例
    net = SfSNet()
    # 載入參數
    net.load_weights_from_pkl('SfSNet-Caffe/weights.pkl')
    # 設置爲測試模式
    net.eval()
    # 讀取並預處理圖像
    image = read_image('data/1.png_face.png')
    # 前向傳播
    out = net(torch.from_numpy(image))
    # 保存
    np.save('conv3.pytorch.npy', out[0].detach().numpy())

• 比對結果：

if __name__ == '__main__':
    caffe_result = np.load('conv3.caffe.npy')
    torch_result = np.load('conv3.pytorch.npy')
    # same函數之前有提到
    print(same(caffe_result, torch_result))

3、比對策略

可以一層一層地比對，也可以抽出中間的某一層來比對，確定不正確的範圍。

4、注意事項

一定要結合可視化結果來比對啊，因爲Caffe模型裏有很多“In-Place”操作，後面的層運行結束之後，會把前面的層的值修改了！！！ 舉個例子：

從可視化結果來看，lconcat1的輸入是nrelu6r和arelu6r，然而，在.prototxt文件裏:

layer {
    name: "lconcat1"
    bottom: "nsum5"
    bottom: "asum5"
    top: "lconcat1"
    type: "Concat"
    concat_param {
        axis: 1
    }
}

lconcat1的輸入是nsum5和asum5。在.prototxt這樣寫確實是對的，因爲nbn6r/abn6r/nrelu6r/arelu6r全是“In-Place層”，修改過的結果又存在nsum5/asum5裏面去了。