博客

近来在训练检测网络的时候会出现loss为nan的情况，需要中断重新训练，会很麻烦。因而选择使用PyTorch提供的梯度裁剪库来对模型训练过程中的梯度范围进行限制，修改之后，不再出现loss为nan的情况。

PyTorch中采用torch.nn.utils.clip_grad_norm_来实现梯度裁剪，链接如下：

https://pytorch.org/docs/stable/_modules/torch/nn/utils/clip_grad.html

训练代码使用示例如下：

from torch.nn.utils import clip_grad_norm_

outputs = model(data)
loss= loss_fn(outputs, target)
optimizer.zero_grad()
loss.backward()

# clip the grad
clip_grad_norm_(model.parameters(), max_norm=20, norm_type=2)

optimizer.step()

其中，max_norm为梯度的最大范数，也是梯度裁剪时主要设置的参数。

备注：网上有同学提醒在（强化学习）使用了梯度裁剪之后训练时间会大大增加。目前在我的检测网络训练中暂时还没有碰到这个问题，以后遇到再来更新。