torch.nan_to_num 是 PyTorch 中用于将张量中的 NaN、正无穷、负无穷替换为指定数值的函数。它的主要作用是清理数据中的异常值,避免它们在后续计算中传播(比如 NaN 会污染整个梯度或损失)。
函数签名
python
torch.nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None)
参数说明
| 参数 | 含义 | 默认行为 |
|---|---|---|
input |
输入张量 | --- |
nan |
替换 NaN 的值 | 0.0 |
posinf |
替换 +∞ 的值 | 若为 None,则用输入 dtype 能表示的最大有限值 |
neginf |
替换 -∞ 的值 | 若为 None,则用输入 dtype 能表示的最小有限值 |
out |
输出张量 | --- |
基本用法
python
import torch
x = torch.tensor([1.0, float('nan'), float('inf'), float('-inf'), 2.0])
# 默认替换
y = torch.nan_to_num(x)
# tensor([1., 0., 3.4028e+38, -3.4028e+38, 2.])
# 自定义替换值
y = torch.nan_to_num(x, nan=0.0, posinf=1e6, neginf=-1e6)
# tensor([1., 0., 1.0e+06, -1.0e+06, 2.])
直观例子
batch_size=2
| 图片 | weather | timeofday | scene |
|---|---|---|---|
| 图 1 | clear(0) | -1(undefined) | highway(0) |
| 图 2 | rain(1) | -1(undefined) | residential(1) |
这个 batch 中 timeofday 全部是 -1
- weather、scene:正常计算 loss
- timeofday:全部被 ignore → CrossEntropyLoss 输出
nan - 总和
val_loss = loss_weather + nan + loss_scene→ 结果 =nan
典型使用场景
损失函数中出现 NaN/Inf 时:在计算 loss 前对预测值或中间结果做清理,防止训练崩溃。
python
loss = criterion(torch.nan_to_num(pred), target)
归一化/除法后:分母为 0 会产生 Inf 或 NaN,可用它兜底。
python
x = a / b
x = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
混合精度训练:fp16 容易溢出成 Inf,用它做保护。
数据预处理:清洗含缺失值(NaN)的特征。