pytorch打印模型梯度

简介

有时候在调试模型训练过程时，我们需要打印模型中参数的梯度，去查看是否存在梯度消失或者梯度爆炸的问题。可以通过在backward之后查看params的grad属性来确认。

参考代码如下所示。

import torch

# 定义模型
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, 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.relu(x)
        x = self.fc2(x)
        x = torch.softmax(x, dim=1)
        return x

# 创建模型实例
model = MyModel()

# 定义输入和目标向量
x = torch.randn(2, 10)
y = torch.tensor([[0, 1], [1, 0]], dtype=torch.float32)

# 进行前向传递
out = model(x)

# 计算损失
loss = torch.nn.functional.binary_cross_entropy(out, y)

# 根据损失计算梯度
loss.backward()

# 收集每个可学习参数的梯度
grads = {}
for name, param in model.named_parameters():
    if param.requires_grad and param.grad is not None:
        grads[name] = param.grad

# 输出梯度
print(grads)