Pytorch使用钩子函数将变量梯度为 nan 的修改为 1

原创已于 2023-12-20 14:44:57 修改 · 485 阅读

CC 4.0 BY-SA版权

文章标签：

于 2023-12-20 14:09:23 首次发布

8 篇文章

订阅专栏

本文解释了在PyTorch中，为了处理计算过程中可能出现的NaN梯度，如何在`backward()`函数之前使用`torch.where()`和自定义hook函数`lambdagrad`。

需要注意一点，代码需要写在 backward() 前面，因为钩子函数是在 backward 时执行的。

            var.register_hook(lambda grad: torch.where(torch.isnan(grad), torch.ones_like(grad), grad))
         
            total_loss.backward()

Pytorch使用 钩子函数 将变量梯度 为 nan 的修改为 1