加載模型字典

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

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

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()

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

最近在做目標檢測，需要把訓好的模型放到嵌入式設備上跑前向，因此得把各種層的實現都用C手擼一遍，，，此為背景。

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

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

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

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

最新評論