如何选取两个torch张量中非-100的对应值并计算MSELoss

你可以通过布尔掩码筛选有效位置元素后再计算损失，完整实现代码如下：

import torch
import torch.nn as nn

# 定义输入张量
predictions = torch.tensor([[33, 34,  7,  5,  5, 23, 22,  1,  3,  5, 23,  1],
        [14,  1, 22,  7,  5, 11,  7, 33,  3, 12, 25, 22],
        [33,  1, 14, 12, 23, 22, 12,  2,  3, 12, 23, 14],
        [23, 34, 34,  3,  5, 25, 12, 11,  2, 23, 23, 13]])

labels = torch.tensor([[-100, -100, -100, -100, -100,   11, -100, -100, -100, -100, -100, -100],
        [-100, -100, -100, -100, -100, -100, -100, -100,   40, -100, -100, -100],
        [-100,   42, -100,   43, -100, -100, -100, -100, -100, -100, -100, -100],
        [-100, -100, -100, -100, -100, -100, -100, -100, -100, -100,   32, -100]])

# 生成有效位置掩码：标签不等于-100的位置为True，其余为False
mask = labels != -100

# 筛选有效元素，转为float避免整数计算的精度问题
valid_preds = predictions[mask].float()
valid_labels = labels[mask].float()

# 计算MSE损失
loss_fn = nn.MSELoss()
loss = loss_fn(valid_preds, valid_labels)

print(loss) # 输出结果为 tensor(847.2000)

步骤说明 #

• 掩码生成：labels != -100会生成和输入张量形状完全一致的布尔张量，自动标记所有需要参与计算的位置
• 元素筛选：用布尔掩码直接索引张量，会提取所有True位置的元素，返回一维张量，自动保持预测值和标签的对应关系
• 损失计算：直接对筛选后的一维张量计算MSE即可，结果仅包含有效位置的误差

你也可以手动实现损失计算，和调用nn.MSELoss结果完全一致：

loss = torch.mean((valid_preds - valid_labels) ** 2)

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