torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
时间: 2024-09-07 13:00:46 浏览: 163
`torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)` 是PyTorch库中用于梯度裁剪(gradient clipping)的一个实用工具。这个函数主要用于防止模型训练过程中由于梯度过大导致的数值不稳定或爆炸性增长。它接受两个参数:
1. `model.parameters()`:这是对模型所有参数的迭代器,`clip_grad_norm_`会遍历这些参数并应用裁剪操作。
2. `max_norm`:设置了一个阈值,即每个参数组(如权重矩阵)的梯度范数不能超过这个值。一旦某个参数组的梯度范数超过了`max_norm`,则会按比例缩放其梯度,使其保持在指定的界限内。
举个简单的例子[^1]:
```python
# 假设 model 是一个有参数的神经网络模型
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
for inputs, targets in dataloader:
optimizer.zero_grad() # 梯度清零
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward() # 反向传播计算梯度
# 使用clip_grad_norm_进行梯度裁剪
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # 设置最大允许的梯度范数为1.0
optimizer.step() # 更新模型参数
```
阅读全文