pytorch钩子函数

一、定义

1. 定义
2. 案例1-分析某一层结果
3. 案例2-挂载到模型上，分析每一层参数
4. 案例3-挂载到fc层反向传播上，增加噪音
5. 模块可视化

二、实现

1. 定义
   钩子函数是指在程序执行过程中预留的一些接口，可以在特定的时间点插入自定义的代码，从而改变程序的行为。在PyTorch中，钩子函数可以在模型训练过程中捕获中间结果，对其进行分析或者可视化，也可以在模型的参数更新过程中进行一些自定义操作。
   
   参考：https://zhuanlan.zhihu.com/p/628354472

2. 案例1-分析某一层结果

import torch
import torch.nn as nn
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 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 = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, (2, 2))
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = x.view(-1, self.num_flat_features(x))
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

def hook_fn(module, input, output):
    for name, param in module.named_parameters():
        if param.requires_grad:
            print(f"Parameter {name}: {param.data}")
    print('input:', input)
    print('output:', output)

net = Net()
net.conv1.register_forward_hook(hook_fn)   #执行完conv1 方法即可执行该钩子

input = torch.randn(1, 3, 32, 32)
output = net(input)

3. 案例2-挂载到模型上，分析每一层参数

import torch
# 定义一个钩子函数，用于记录模型参数
def record_param(module, input, output):
    for name, param in module.named_parameters():
        if param.requires_grad:
            print(f"Parameter {name}: {param.data}")
# 定义一个简单的模型
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = torch.nn.Linear(10, 5)
        self.fc2 = torch.nn.Linear(5, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.nn.functional.relu(x)
        x = self.fc2(x)
        return x

# 实例化模型和数据
net = Net()
data = torch.randn(1, 10)
# 注册钩子函数
handle = net.register_forward_hook(record_param)
# 前向传播
output = net(data)
# 移除钩子函数
handle.remove()

4. 案例3-挂载到fc层反向传播上，增加噪音

import torch

# 定义一个钩子函数，用于修改梯度
def add_noise_grad(grad):
    noise = torch.randn_like(grad) * 0.01
    return grad + noise

# 定义一个简单的模型
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = torch.nn.Linear(10, 5)
        self.fc2 = torch.nn.Linear(5, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.nn.functional.relu(x)
        x = self.fc2(x)
        return x

# 实例化模型和数据
net = Net()
data = torch.randn(1, 10)
label = torch.tensor([0, 1])

# 注册钩子函数
handle = net.fc2.register_backward_hook(add_noise_grad)

# 前向传播和反向传播
output = net(data)
loss = torch.nn.functional.cross_entropy(output, label)
loss.backward()

# 移除钩子函数
handle.remove()

5. 模块可视化

import torch
from torch import nn
from torchviz import make_dot

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x

net = Net()

def hook_fn(m, i, o):
    print(m)
    print("------------Input Grad------------")
    print(i)
    print("------------Output Grad-----------")
    print(o)

handle = net.fc1.register_forward_hook(hook_fn)

x = torch.randn(1, 10)
y = net(x)
z=make_dot(y, params=dict(net.named_parameters()))
z.view()