Torch 中实现掩码加权

1. PyTorch 中实现掩码加权的 L1 损失计算（用于图像配准或相似度评估）

（单张图像）

import torch

def masked_l1_loss(pred, target, mask, eps=1e-6):
    """
    pred:   预测图像 [C, H, W] 或 [D, H, W]
    target: 目标图像 [C, H, W] 或 [D, H, W]
    mask:   掩码 [1, H, W] 或 [D, H, W]（与target同空间维度）
    """
    # 计算逐像素L1距离
    l1 = torch.abs(pred - target)
    
    # 掩码加权求和
    masked_l1 = l1 * mask
    
    # 归一化（有效区域的平均L1）
    loss = masked_l1.sum() / (mask.sum() + eps)
    
    return loss

(2) 批量处理版本（支持 Batch 维度）

def masked_l1_loss_batch(pred, target, mask, eps=1e-6):
    """
    pred:   [B, C, H, W] 或 [B, D, H, W]
    target: [B, C, H, W] 或 [B, D, H, W]
    mask:   [B, 1, H, W] 或 [B, D, H, W]
    """
    l1 = torch.abs(pred - target)
    masked_l1 = l1 * mask
    
    # 按Batch求平均（每个样本独立归一化）
    loss_per_sample = masked_l1.sum(dim=(1,2,3)) / (mask.sum(dim=(1,2,3)) + eps)
    return loss_per_sample.mean()

(4) 模拟数据生成

# 创建示例数据（3D图像）
B, D, H, W = 2, 64, 128, 128
pred = torch.randn(B, 1, D, H, W)  # 预测图像
target = torch.randn(B, 1, D, H, W) # 目标图像

# 生成随机掩码（模拟重叠区域）
mask = (torch.rand(B, 1, D, H, W) > 0.3).float()  # 70%区域有效

(5) 计算损失

loss = masked_l1_loss_batch(pred, target, mask)
print(f"Masked L1 Loss: {loss.item():.4f}")

(6) 梯度回传（用于训练）

loss.backward() 

---

6. 关键点说明

操作	作用
torch.abs()	计算 L1 距离
* mask	掩码加权（无效区域贡献为 0）
sum() / (sum + eps)	仅在有效区域内归一化，避免背景区域干扰
GPU 加速	所有操作支持 CUDA，适合大规模医学图像

7. 扩展功能

(1) 多通道图像处理

若图像有多个通道（如 RGB 或多模态），需调整掩码维度：

mask = mask.expand_as(target)  # 将[B,1,H,W]扩展为[B,C,H,W]

(2) 加权掩码（非二进制）

如果掩码是软掩码（如置信度权重）：

soft_mask = torch.sigmoid(weight_map)  # 值域[0,1]
loss = (torch.abs(pred - target) * soft_mask).sum() / soft_mask.sum()

(3) 与 SmoothL1 结合

def masked_smooth_l1(pred, target, mask, beta=1.0):
    diff = torch.abs(pred - target)
    smooth_l1 = torch.where(diff < beta, 0.5 * diff**2 / beta, diff - 0.5 * beta)
    return (smooth_l1 * mask).sum() / (mask.sum() + 1e-6)

通过这种方式，可以确保 L1 损失仅评估图像配准中真正有意义的解剖结构重叠区域，显著提升配准精度和训练稳定性。