使用掩码时的MSELoss

import torch
import torch.nn as nn
import torch.nn.functional as F

class MaskedMSELoss(nn.Module):
    def __init__(self):
        super(MaskedMSELoss, self).__init__()

    def forward(self, pred, target, mask):
        # 计算预测值与真实值之间的平方差
        squared_diff = (pred - target) ** 2
        # 应用掩码
        masked_squared_diff = squared_diff * mask
        # 计算平均损失
        loss = masked_squared_diff.sum() / mask.sum()
        return loss

# 示例数据
pred = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
target = torch.tensor([[1.5, 2.5, 3.5], [4.5, 5.5, 6.5]])
mask = torch.tensor([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]])

# 创建损失函数实例
criterion = MaskedMSELoss()

# 计算损失
loss = criterion(pred, target, mask)
print(f"Masked MSE Loss: {loss.item()}")