将pytorch中变量的梯度为 nan 的替换成 1 还是 0？

培之

于 2024-02-19 14:50:28 发布

阅读量739

点赞数 11

文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/OrdinaryMatthew/article/details/136170365

版权

本文介绍了如何在PyTorch中使用xtensor创建张量，执行基本运算，如乘法和立方，并演示了如何应用梯度替换技巧来控制优化过程。作者展示了如何使用Adam优化器对张量进行学习并观察梯度替换对模型更新的影响。

摘要由CSDN通过智能技术生成

替换成0，则变量保持不动

0: xtensor([1.0000, 2.0000, 3.0000, 4.5000], device='cuda:0', requires_grad=True)
0: xtensor([1.0000, 2.0000, 3.0000, 4.5000], device='cuda:0', requires_grad=True)

替换成1，变量会变化

0: xtensor([1.0000, 2.0000, 3.0000, 4.5000], device='cuda:0', requires_grad=True)
0: xtensor([1.0000, 2.0000, 3.0000, 4.5000], device='cuda:0', requires_grad=True)

代码如下：

import torch
import numpy as np
import torch.optim as optim

torch.autograd.set_detect_anomaly(True)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
x = torch.tensor([1.0, 2.0, 3.0,4.5], dtype=torch.float32, requires_grad=True, device=device)

y1 = x*2
y2 = x**3
print(f'x{x}')
print(f'y1{y1}')
print(f'y2{y2}')

optimizer = optim.Adam([x], lr=1)

def replace_gradients_zero(grad):  
    return torch.where(torch.ones_like(grad, dtype=torch.bool), torch.zeros_like(grad), grad)  

def replace_gradients_one(grad):  
    return torch.where(torch.ones_like(grad, dtype=torch.bool), torch.ones_like(grad), grad)  

for i in range(1):
    print(f'{i}: x{x}')
    optimizer.zero_grad()
    loss = (y1-y2).sum()

    for param in [x]:  
        param.register_hook(replace_gradients_zero) 

    loss.backward()
    optimizer.step()
    print(f'{i}: x{x}')