『PyTorch』第十六彈_hook技術

由於pytorch會自動舍棄圖計算的中間結果，所以想要獲取這些數值就需要使用鈎子函數。

鈎子函數包括Variable的鈎子和nn.Module鈎子，用法相似。

 

一、register_hook

import torch
from torch.autograd import Variable

grad_list = []

def print_grad(grad):
    grad_list.append(grad)

x = Variable(torch.randn(2, 1), requires_grad=True)
y = x+2
z = torch.mean(torch.pow(y, 2))
lr = 1e-3
y.register_hook(print_grad)
z.backward()
x.data -= lr*x.grad.data

print(grad_list)

[Variable containing:
 1.5653
 3.5175
[torch.FloatTensor of size 2x1]
]

 

二、register_forward_hook & register_backward_hook

這兩個函數的功能類似於variable函數的register_hook，可在module前向傳播或反向傳播時注冊鈎子。

每次前向傳播執行結束后會執行鈎子函數（hook）。前向傳播的鈎子函數具有如下形式：hook(module, input, output) -> None，而反向傳播則具有如下形式：hook(module, grad_input, grad_output) -> Tensor or None。

鈎子函數不應修改輸入和輸出，並且在使用后應及時刪除，以避免每次都運行鈎子增加運行負載。鈎子函數主要用在獲取某些中間結果的情景，如中間某一層的輸出或某一層的梯度。這些結果本應寫在forward函數中，但如果在forward函數中專門加上這些處理，可能會使處理邏輯比較復雜，這時候使用鈎子技術就更合適一些。下面考慮一種場景，有一個預訓練好的模型，需要提取模型的某一層（不是最后一層）的輸出作為特征進行分類，但又不希望修改其原有的模型定義文件，這時就可以利用鈎子函數。下面給出實現的偽代碼。

model = VGG() features = t.Tensor() def hook(module, input, output): '''把這層的輸出拷貝到features中''' features.copy_(output.data) handle = model.layer8.register_forward_hook(hook) _ = model(input) # 用完hook后刪除 handle.remove()

測試LeNet網絡

import torch as t
import torch.nn as nn
import torch.nn.functional as F
  
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet,self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6,16,5)
        self.fc1 = nn.Linear(16*5*5,120)
        self.fc2 = nn.Linear(120,84)
        self.fc3 = nn.Linear(84,10)
  
    def forward(self,x):
        x = F.max_pool2d(F.relu(self.conv1(x)),(2,2))
        x = F.max_pool2d(F.relu(self.conv2(x)),2)
        x = x.view(x.size()[0], -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

 先模擬一下單次的向前傳播，

net = LeNet()
img = t.autograd.Variable((t.arange(32*32*1).view(1,1,32,32)))
net(img)

Variable containing:

Columns 0 to 7 
 27.6373 -13.4590  23.0988 -16.4491  -8.8454 -15.6934  -4.8512   1.3490

Columns 8 to 9 
  3.7801 -15.9396
[torch.FloatTensor of size 1x10]

仿照上面示意，進行鈎子注冊，獲取第一卷積層輸出結果，

def hook(module, inputdata, output):
    '''把這層的輸出拷貝到features中'''
    print(output.data)

handle = net.conv2.register_forward_hook(hook)
net(img)
# 用完hook后刪除
handle.remove()

……

……

 [torch.FloatTensor of size 1x16x10x10]

 看看hook能識別什么

import torch
from torch import nn
import torch.functional as F
from torch.autograd import Variable

def for_hook(module, input, output):
    print(module)
    for val in input:
        print("input val:",val)
    for out_val in output:
        print("output val:", out_val)

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
    def forward(self, x):

        return x+1

model = Model()
x = Variable(torch.FloatTensor([1]), requires_grad=True)
handle = model.register_forward_hook(for_hook)
print(model(x))
handle.remove()

 可見對於目標層，其輸入輸出都可以獲取到，

Model(
)
input val: Variable containing:
           1
          [torch.FloatTensor of size 1]

output val: Variable containing:
            2
           [torch.FloatTensor of size 1]

Variable containing:
 2
[torch.FloatTensor of size 1]