PyTorch中如何跳过含NaN无效样本的损失计算与梯度回传

问题背景 #

• 任务基于自编码网络实现姿态条件人脸生成，输入为原始图像+姿态条件向量，生成结果受姿态条件约束。
• 训练阶段采用MSE作为关键点重建损失时，训练初期网络输出混乱无意义，人脸关键点检测库会对无效输出返回None，无法得到预期形状为N x 3 x 256 x 256的关键点张量。
• 现有处理逻辑为对无效关键点样本填充NaN，要求MSE计算时自动跳过全NaN的无效样本，不参与梯度回传，替代低效的Python层for循环逐样本计算方案。

实现逻辑 #

全程使用PyTorch向量化运算，无逐样本循环，性能与原生MSE损失一致：

1. 逐样本生成有效掩码：判断每个样本是否全为NaN，全NaN标记为无效样本，其余为有效样本
2. 计算逐像素平方误差，将NaN位置的误差值置0，避免数值异常传播
3. 通过掩码屏蔽无效样本的所有误差值，仅对有效样本计算平均损失，无效样本不参与损失统计与梯度回传
4. 兼容批次内全为无效样本的边界场景，避免除0报错

完整可运行代码 #

import numpy as np
import torch
import random
from torchvision import transforms

# 初始化图像预处理
size = 256
transform = transforms.Compose(
    [
        transforms.ToPILImage(),
        transforms.CenterCrop(size),
        transforms.ToTensor(),
    ]
)

# 模拟关键点生成函数
def generate_mesh_from_image(image_array):
    rand = random.randint(0, 1)
    if rand == 0:
        return None
    else:
        return np.random.randn(3, 256, 256).astype(np.float32)

def build_landmark_batch(image_tensors):
    tensor_meshes = []
    for image_tensor in image_tensors:
        image_array = image_tensor.detach().cpu().numpy()
        image_landmarks = generate_mesh_from_image(image_array)
        if image_landmarks is None:
            # 修正原代码np.empty传参错误，显式为无效样本填充NaN
            image_landmarks = np.full((3, 256, 256), np.nan, dtype=np.float32)
        landmark_tensor = transform(image_landmarks)
        tensor_meshes.append(landmark_tensor)
    # 拼接为批次张量，输出形状 (batch_size, 3, 256, 256)
    return torch.stack(tensor_meshes, dim=0)

def masked_mse_loss(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """
    自动忽略全NaN样本的MSE损失，支持反向传播
    Args:
        pred: 预测关键点张量，形状 (B, C, H, W)，无效样本位置填充NaN
        target: 真实关键点张量，形状与pred一致
    """
    # 生成有效样本掩码：形状(B,)，有效样本对应值为True
    valid_sample_mask = ~torch.isnan(pred).flatten(start_dim=1).all(dim=1)
    valid_count = valid_sample_mask.sum()
    # 处理批次全无效的边界情况，返回零梯度张量避免训练中断
    if valid_count == 0:
        return pred.sum() * 0.0
    # 计算逐像素平方误差，NaN位置置0
    square_error = (pred - target) ** 2
    square_error = torch.nan_to_num(square_error, nan=0.0)
    # 扩展掩码形状匹配误差张量，屏蔽无效样本的误差
    mask_expanded = valid_sample_mask.view(-1, 1, 1, 1)
    masked_error = square_error * mask_expanded
    # 仅对有效样本计算平均MSE
    total_elements = valid_count * pred.shape[1] * pred.shape[2] * pred.shape[3]
    loss = masked_error.sum() / total_elements
    return loss

# 功能测试
if __name__ == "__main__":
    image_tensors = torch.randn(10, 3, 256, 256)
    pred_landmarks = build_landmark_batch(image_tensors)
    gt_landmarks = torch.randn_like(pred_landmarks)
    loss = masked_mse_loss(pred_landmarks, gt_landmarks)
    loss.backward() # 反向传播正常执行，无效样本无梯度贡献

注意事项 #

• 无效样本必须保证所有像素位置均填充为NaN，若样本内存在非NaN值会被判定为有效样本参与损失计算
• 该实现所有运算均部署在PyTorch计算图内，支持GPU加速、自动混合精度训练，无设备兼容问题
• 若需要调整为逐样本加权损失，只需在masked_error求和前乘以对应权重即可，逻辑无需改动

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