0、參考

https://oldpan.me/archives/pytorch-autograd-hook

https://pytorch.org/docs/stable/search.html?q=hook&check_keywords=yes&area=default

https://github.com/pytorch/pytorch/issues/598

https://github.com/sksq96/pytorch-summary

https://github.com/allensll/test/blob/591c7ce3671dbd9687b3e84e1628492f24116dd9/net_analysis/viz_lenet.py

1、背景

在神經網絡的反向傳播當中個，流程只保存葉子節點的梯度，對于中間變量的梯度沒有進行保存。

import torch x = torch.tensor([1,2],dtype=torch.float32,requires_grad=True) y = x+2 z = torch.mean(torch.pow(y, 2)) lr = 1e-3 z.backward() x.data -= lr*x.grad.data print(y.grad)

此時輸出就是：None，這個時候hook的作用就派上，hook可以通過自定義一些函數，從而完成中間變量的輸出，比如中間特征圖、中間層梯度修正等。

? 在pytorch docs搜索hook，可以發現有四個hook相關的函數，分別為register_hook，register_backward_hook，register_forward_hook，register_forward_pre_hook。其中register_hook屬于tensor類，而后面三個屬于moudule類。

• register_hook函數屬于torch.tensor類，函數在tensor梯度計算的時候就會執行，這個函數主要處理梯度相關的數據，表現形式$hook(grad) rightarrow Tensor or None$.

import torch x = torch.tensor([1,2],dtype=torch.float32,requires_grad=True) y = x * 2 y.register_hook(print) <torch.utils.hooks.RemovableHandle at 0x7f765e876f60> z = torch.mean(y) z.backward() tensor([ 0.5000, 0.5000])

• Register_backward_hook等三個屬于torch.nn，屬于moudule中的方法。

hook(module, grad_input, grad_output) -> Tensor or None

寫個demo，參考:

下面的計算為

import torch import torch.nn as nn device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")def print_hook(grad):print ("register hook:", grad)return gradclass TestNet(nn.Module):def __init__(self):super(TestNet, self).__init__()self.f1 = nn.Linear(4, 1, bias=True)self.weights_init()def weights_init(self):self.f1.weight.data.fill_(4)self.f1.bias.data.fill_(0.1)def forward(self, input):self.input = inputout = input * 0.75out = self.f1(out)out = out / 4return outdef back_hook(self, moudle, grad_input, grad_output):print ("back hook in:", grad_input)print ("back hook out:", grad_output)# 修改梯度# grad_input = list(grad_input)# grad_input[0] = grad_input[0] * 100# print (grad_input)return tuple(grad_input)if __name__ == '__main__':input = torch.tensor([1, 2, 3, 4], dtype=torch.float32, requires_grad=True).to(device)net = TestNet()net.to(device)net.register_backward_hook(net.back_hook)ret = net(input)print ("result", ret)ret.backward()print('input.grad:', input.grad)for param in net.parameters():print('{}:grad->{}'.format(param, param.grad))

輸出：

result tensor([7.5250], grad_fn=<DivBackward0>) back hook in: (tensor([0.2500]), None) back hook out: (tensor([1.]),) input.grad: tensor([0.7500, 0.7500, 0.7500, 0.7500]) Parameter containing: tensor([[4., 4., 4., 4.]], requires_grad=True):grad->tensor([[0.1875, 0.3750, 0.5625, 0.7500]]) Parameter containing: tensor([0.1000], requires_grad=True):grad->tensor([0.2500])

輸出結果以及梯度都很明顯，簡單分析一下w權重的梯度，

另外，hook中有個bug，假設我們bug，假設我們注釋掉out = out / 4這行，可以發現輸出變成back hook in: (tensor([1.]), tensor([1.]))。這種情況就不符合上面我們的梯度計算公式，是因為這個時候：

則此時的偏導只是對

和 進行計算，所以都是1，1。這是pytorch的設計缺陷

• register_forward_hook跟Register_backward_hook差不多，就不過多復述。
• register_forward_pre_hook，可以發現其輸入只有hook(module, input) -> None
  其主要是針對推理時的hook.

2、應用

2.1 特征圖打印

? 直接利用pytorch已有的resnet18進行特征圖打印，只打印卷積層的特征圖，

import torch from torchvision.models import resnet18 import torch.nn as nn from torchvision import transformsimport matplotlib.pyplot as pltdef viz(module, input):x = input[0][0]#最多顯示4張圖min_num = np.minimum(4, x.size()[0])for i in range(min_num):plt.subplot(1, 4, i+1)plt.imshow(x[i].cpu())plt.show()import cv2 import numpy as np def main():t = transforms.Compose([transforms.ToPILImage(),transforms.Resize((224, 224)),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])])device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")model = resnet18(pretrained=True).to(device)for name, m in model.named_modules():# if not isinstance(m, torch.nn.ModuleList) and # not isinstance(m, torch.nn.Sequential) and # type(m) in torch.nn.__dict__.values():# 這里只對卷積層的feature map進行顯示if isinstance(m, torch.nn.Conv2d):m.register_forward_pre_hook(viz)img = cv2.imread('./cat.jpeg')img = t(img).unsqueeze(0).to(device)with torch.no_grad():model(img)if __name__ == '__main__':main()

直接放幾張中間層的圖

圖1 第一層卷積層輸入

圖2 第四層卷積層的輸入

2.2 模型大小，算力計算

同樣的用法，可以直接參考pytorch-summary這個項目。

總結

以上是生活随笔為你收集整理的pytorch forward_pytorch使用hook打印中间特征图、计算网络算力等的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。