PyTorch Tensor赋值时值变为0而非指定计算值的异常问题

复现代码 #

import torch
def normalize(x, x_min, x_max):
       return x/(x_max - x_min)

x_min = torch.tensor([-5,  0])
x_max = torch.tensor([10, 15])

old = torch.tensor([[-3,-3],[3,3]])
for i in range(old.shape[0]):
    old[i] = normalize(old[i], x_min, x_max)
print(old)

上述代码执行后输出全0张量，和预期的归一化浮点结果不符。

问题原因 #

• 初始化old张量时传入的是整数列表，未显式指定数据类型，PyTorch会默认创建torch.int64（整数类型）张量
• normalize函数的计算结果是浮点类型张量，将浮点值写入整数类型的张量切片时，PyTorch会做隐式类型转换，直接丢弃小数部分取整。示例中归一化计算结果为[-0.2, -0.2]和[0.2, 0.2]，绝对值均小于1，取整后数值全部为0。

修复方法 #

归一化计算本身需要保留浮点精度，最直接的修复方式是初始化原张量时指定为浮点类型：

import torch
def normalize(x, x_min, x_max):
    return x/(x_max - x_min)

x_min = torch.tensor([-5,  0])
x_max = torch.tensor([10, 15])

# 初始化时指定dtype为浮点型
old = torch.tensor([[-3,-3],[3,3]], dtype=torch.float32)
for i in range(old.shape[0]):
    old[i] = normalize(old[i], x_min, x_max)
print(old)

执行后得到预期输出：

tensor([[-0.2000, -0.2000],
        [ 0.2000,  0.2000]])

代码优化提示：该场景无需逐行循环赋值，依托PyTorch的广播机制可以直接对整个张量做批量计算，代码更简洁、运行效率更高：

old = torch.tensor([[-3,-3],[3,3]], dtype=torch.float32)
old = normalize(old, x_min, x_max)

内容的提问来源于stack exchange，提问作者Neel Shah