加載模型字典

逐一判斷每一層，如果該層是bn 的 running mean，就取出參數(shù)并取平均作為該層的代表

對保存的每個BN層的數(shù)值進行曲線可視化

from functools import partial
import pickle
import torch
import matplotlib.pyplot as plt

pth_path = 'checkpoint.pth'

pickle.load = partial(pickle.load, encoding="latin1")
pickle.Unpickler = partial(pickle.Unpickler, encoding="latin1")
pretrained_dict = torch.load(pth_path, map_location=lambda storage, loc: storage, pickle_module=pickle)
pretrained_dict = pretrained_dict['state_dict']

means = []
for name, param in pretrained_dict.items():
 print(name)
 if 'running_mean' in name:
  means.append(mean.numpy())

layers = [i for i in range(len(means))]

plt.plot(layers, means, color='blue')
plt.legend()
plt.xticks(layers)
plt.xlabel('layers')
plt.show()

補充知識：關(guān)于pytorch中BN層（具體實現(xiàn)）的一些小細節(jié)

最近在做目標(biāo)檢測，需要把訓(xùn)好的模型放到嵌入式設(shè)備上跑前向，因此得把各種層的實現(xiàn)都用C手擼一遍，，，此為背景。

其他層沒什么好說的，但是BN層這有個小坑。pytorch在打印網(wǎng)絡(luò)參數(shù)的時候，只打出weight和bias這兩個參數(shù)。咦，說好的BN層有四個參數(shù)running_mean、running_var 、gamma 、beta的呢？一開始我以為是pytorch把BN層的計算簡化成weight * X + bias，但馬上反應(yīng)過來應(yīng)該沒這么簡單，因為pytorch中只有可學(xué)習(xí)的參數(shù)才稱為parameter。上網(wǎng)找了一些資料但都沒有說到這么細的，畢竟大部分用戶使用時只要模型能跑起來就行了，，，于是開始看BN層有哪些屬性，果然發(fā)現(xiàn)了熟悉的running_mean和running_var，原來pytorch的BN層實現(xiàn)并沒有不同。這里吐個槽：為啥要把gamma和beta改叫weight、bias啊，很有迷惑性的好不好，，，

扯了這么多，干脆捋一遍pytorch里BN層的具體實現(xiàn)過程，幫自己理清思路，也可以給大家提供參考。再吐槽一下，在網(wǎng)上搜“pytorch bn層”出來的全是關(guān)于這一層怎么用的、初始化時要輸入哪些參數(shù)，沒找到一個pytorch中BN層是怎么實現(xiàn)的，，，

眾所周知，BN層的輸出Y與輸入X之間的關(guān)系是：Y = (X - running_mean) / sqrt(running_var + eps) * gamma + beta，此不贅言。其中g(shù)amma、beta為可學(xué)習(xí)參數(shù)（在pytorch中分別改叫weight和bias），訓(xùn)練時通過反向傳播更新；而running_mean、running_var則是在前向時先由X計算出mean和var，再由mean和var以動量momentum來更新running_mean和running_var。所以在訓(xùn)練階段，running_mean和running_var在每次前向時更新一次；在測試階段，則通過net.eval()固定該BN層的running_mean和running_var，此時這兩個值即為訓(xùn)練階段最后一次前向時確定的值，并在整個測試階段保持不變。

以上這篇可視化pytorch 模型中不同BN層的running mean曲線實例就是小編分享給大家的全部內(nèi)容了，希望能給大家一個參考，也希望大家多多支持腳本之家。

最新評論