如何自定义卷积神经网络损失函数，仅计算张量指定区域以忽略边框损失

方法1：生成同步掩码过滤边框（推荐，适配所有预处理逻辑） #

这种方法不需要计算缩放比例，不会因为预处理参数调整出错，只需要修改DataLoader和损失计算逻辑：

第一步：修改DataLoader的__getitem__方法，同步返回图像和掩码 #

def __getitem__(self, index):
    path = self.input_data[index]
    imgs_path = sorted(glob.glob(path + '/*.png'))
    #read light conditions
    lightConditions = []
    with open(path + "/lightConditions.json", 'r') as file:
        lightConditions = json.load(file)
    #shift light conditions
    lightConditions.pop(0)
    lightConditions.append(False)
    frameNumber = 0
    imgs = []
    masks = [] # 新增掩码存储列表
    for img_path in imgs_path:
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        im_pil = Image.fromarray(img)
        # 生成原始无框区域的掩码：无框区域为1，后续加边框的位置为0
        mask_pil = Image.new('L', im_pil.size, 1)
        if lightConditions[frameNumber] ==False:
            imgBorder = ImageOps.expand(im_pil,border = 6, fill='black')
            maskBorder = ImageOps.expand(mask_pil, border=6, fill=0) # 边框位置掩码设为0
        else:
            imgBorder = ImageOps.expand(im_pil, border = 6, fill='orange')
            maskBorder = ImageOps.expand(mask_pil, border=6, fill=0) # 边框位置掩码设为0
        img = np.asarray(imgBorder)
        mask = np.asarray(maskBorder)
        # 图像和掩码做完全相同的resize操作
        img = cv2.resize(img, (256,448))
        mask = cv2.resize(mask, (256,448), interpolation=cv2.INTER_NEAREST) # 掩码用最近邻插值避免出现非0/1值
        img_tensor = ToTensor()(img).float()
        mask_tensor = ToTensor()(mask).float()
        imgs.append(img_tensor)
        masks.append(mask_tensor)
        frameNumber +=1
    imgs = torch.stack(imgs, dim=0)
    masks = torch.stack(masks, dim=0)
    return imgs, masks # 同时返回图像和掩码

第二步：修改训练逻辑，加载对应GT的掩码 #

for index, (data_sr, mask_sr) in enumerate(train_loader): # 同步加载图像和掩码
    data_sr = data_sr.cuda(processGPU, non_blocking=True)
    mask_sr = mask_sr.cuda(processGPU, non_blocking=True)
    # 原有逻辑不变
    data_sr = data_sr.squeeze()
    mask_sr = mask_sr.squeeze()
    num_f = len(data_sr)
    if cfg.model_type == 0:
        idx_start = random.randint(-2, 2)
        idx_all = list(np.arange(idx_start, idx_start + num_f).clip(0, num_f - 1))
        idx_gt = [idx_all.pop(int(num_f / 2))]
        idx_input = idx_all
    else:
        idx_all = np.arange(0, num_f)
        idx_input = list(idx_all[0:4])
        idx_gt = list(idx_all[4:4+cfg.dec_frames])
    imgs_input = data_sr[idx_input]
    imgs_gt = data_sr[idx_gt]
    mask_gt = mask_sr[idx_gt] # 取出GT对应的掩码
    imgs_pred = model(imgs_input)

第三步：自定义损失计算逻辑，仅计算有效区域 #

# 方法A：直接乘掩码后计算MSE
loss = criterion_mse(imgs_pred * mask_gt, imgs_gt * mask_gt)

# 方法B：仅选择有效区域的像素计算，避免边框区域的0值影响平均损失
def masked_mse_loss(pred, gt, mask):
    mask = mask.expand_as(pred) # 把单通道掩码扩展到和图像相同的通道数
    valid_pixels = mask.sum()
    if valid_pixels == 0:
        return torch.tensor(0.0, device=pred.device)
    return torch.sum((pred - gt) ** 2 * mask) / valid_pixels

# 调用自定义损失
loss = masked_mse_loss(imgs_pred, imgs_gt, mask_gt)

方法2：直接切片有效区域（适合预处理参数固定的场景） #

如果你所有原始无框图像的尺寸固定，可以直接计算出resize后边框的厚度，直接切片中间有效区域计算损失：

1. 先计算缩放后边框厚度：假设原始无框图像高为H、宽为W，加12px边框后高为H+12、宽为W+12，resize到高256、宽448，则：
   • 上下边框的缩放后厚度：border_h = round(6 * 256 / (H + 12))
   • 左右边框的缩放后厚度：border_w = round(6 * 448 / (W + 12))
2. 直接切片有效区域计算损失：

# 图像张量维度为 [帧数, 通道数, 高, 宽]，高为256，宽为448
border_h = 3 # 替换为你实际计算出的数值
border_w = 3 # 替换为你实际计算出的数值
# 切片中间有效区域
pred_valid = imgs_pred[:, :, border_h:256-border_h, border_w:448-border_w]
gt_valid = imgs_gt[:, :, border_h:256-border_h, border_w:448-border_w]
loss = criterion_mse(pred_valid, gt_valid)

内容的提问来源于stack exchange，提问作者Pia Lüdemann