PyTorch torch.nansum 函数
torch.nansum 是 PyTorch 中用于返回张量的非NaN值之和的函数。与 sum 不同的是,它会忽略所有的 NaN 值。
函数定义
torch.nansum(input, dim, keepdim=False)
使用示例
实例
import torch
x = torch.tensor([1.0, float('nan'), 3.0, 2.0, 5.0])
# 返回非NaN值之和
print("非NaN值之和:", torch.nansum(x))
# 沿 dim=0 非NaN值之和
y = torch.tensor([[1.0, float('nan'), 2.0], [4.0, 1.0, 3.0]])
print("dim=0 非NaN值之和:", torch.nansum(y, dim=0))
print("dim=1 非NaN值之和:", torch.nansum(y, dim=1))
x = torch.tensor([1.0, float('nan'), 3.0, 2.0, 5.0])
# 返回非NaN值之和
print("非NaN值之和:", torch.nansum(x))
# 沿 dim=0 非NaN值之和
y = torch.tensor([[1.0, float('nan'), 2.0], [4.0, 1.0, 3.0]])
print("dim=0 非NaN值之和:", torch.nansum(y, dim=0))
print("dim=1 非NaN值之和:", torch.nansum(y, dim=1))
输出结果为:
非NaN值之和: tensor(11.) dim=0 非NaN值之和: tensor([5., 1., 5.]) dim=1 非NaN值之和: tensor([3., 8.])

Pytorch torch 参考手册