PyTorch梯度裁剪避免训练loss nan的操作

PyTorch梯度裁剪是一种用于避免训练过程中出现loss为nan的问题，其通过限制模型的参数梯度范围来提高训练稳定性和收敛效果。以下是PyTorch梯度裁剪的完整攻略：

什么是梯度裁剪

梯度裁剪是一种通过限制参数梯度范围的方法，防止训练过程中出现梯度爆炸或梯度消失的情况。这种现象常常发生在深层神经网络中，尤其是在使用长短时记忆网络（LSTM）等循环神经网络时更加明显。

常见方法

常见的梯度裁剪方法包括全局范围裁剪和逐层范围裁剪两种。

全局范围裁剪：对所有参数的梯度进行裁剪，即裁剪的范围是所有参数的梯度范围。

逐层范围裁剪：对每个层的参数进行裁剪，裁剪的范围是该层参数的梯度范围。这种方法可更切合实际应用，因为不同层的参数梯度范围差异较大。

操作步骤

PyTorch内置了对梯度裁剪的支持，以下是梯度裁剪的操作步骤：

1. 定义模型，如：

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(100, 10)

    def forward(self, x):
        return self.linear(x)

1. 定义损失函数和优化器

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

1. 对模型进行梯度裁剪

全局裁剪：

clipped_gradients = []
max_grad_norm = 1.0 # 定义最大梯度范围

for group in optimizer.param_groups: # 遍历每个参数组
    for param in group['params']: # 遍历每个参数
        # 计算梯度范数
        grad_norm = param.grad.norm(2)
        # 若梯度范数超过指定范围，则进行梯度裁剪
        if grad_norm > max_grad_norm:
            clipped_gradients.append(param.grad)
            param.grad.div_(grad_norm / max_grad_norm)

逐层裁剪：

clipped_gradients = []
max_grad_norm = 1.0 # 定义最大梯度范围

for group in optimizer.param_groups: # 遍历每个参数组
    for param in group['params']: # 遍历每个参数
        # 若当前参数在线性层中，则计算梯度范数
        if hasattr(param, 'weight'): 
            grad = param.grad
            if grad is None:
                continue
            weight = param.detach()
            if weight.grad is not None:
                weight.grad.detach_()
            else:
                weight.grad = torch.zeros_like(weight)
            # 计算梯度范数
            grad_norm = grad.norm(2)
            # 若梯度范数超过指定范围，则进行梯度裁剪
            if grad_norm > max_grad_norm:
                clipped_gradients.append(grad)
                grad.div_(grad_norm / max_grad_norm)

1. 更新模型参数

optimizer.step()

1. 清空梯度

optimizer.zero_grad()

示例说明

下面举两个示例说明梯度裁剪的操作方法。

示例一：全局梯度裁剪

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(100, 10)

    def forward(self, x):
        return self.linear(x)

model = Model()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# 数据集
x = torch.randn(32, 100)
y = torch.randint(10, (32,))

# 训练
for i in range(100):
    optimizer.zero_grad()
    pred = model(x)
    loss = criterion(pred, y)
    loss.backward()
    # 对梯度进行裁剪
    clipped_gradients = []
    max_grad_norm = 1.0
    for group in optimizer.param_groups:
        for param in group['params']:
            # 计算梯度范数
            grad_norm = param.grad.norm(2)
            # 若梯度范数超过指定范围，则进行梯度裁剪
            if grad_norm > max_grad_norm:
                clipped_gradients.append(param.grad)
                param.grad.div_(grad_norm / max_grad_norm)
    # 更新参数
    optimizer.step()

示例二：逐层梯度裁剪

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(100, 10)
        self.linear2 = nn.Linear(10, 2)

    def forward(self, x):
        x = self.linear1(x)
        x = self.linear2(x)
        return x

model = Model()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# 数据集
x = torch.randn(32, 100)
y = torch.randint(2, (32,))

# 训练
for i in range(100):
    optimizer.zero_grad()
    pred = model(x)
    loss = criterion(pred, y)
    loss.backward()
    # 对梯度进行裁剪
    clipped_gradients = []
    max_grad_norm = 1.0
    for group in optimizer.param_groups:
        for param in group['params']:
            # 若当前参数在线性层中，则计算梯度范数
            if hasattr(param, 'weight'): 
                grad = param.grad
                if grad is None:
                    continue
                weight = param.detach()
                if weight.grad is not None:
                    weight.grad.detach_()
                else:
                    weight.grad = torch.zeros_like(weight)
                # 计算梯度范数
                grad_norm = grad.norm(2)
                # 若梯度范数超过指定范围，则进行梯度裁剪
                if grad_norm > max_grad_norm:
                    clipped_gradients.append(grad)
                    grad.div_(grad_norm / max_grad_norm)
    # 更新参数
    optimizer.step()