Fastai图像分割遇Target越界与CUDA断言错误，求解决方案

我来帮你梳理并解决这两个在Colab环境中遇到的Fastai图像分割问题，结合你使用Fastai 1.0.60做TGS盐识别任务的场景，具体分析如下：

CPU环境错误：Target 255 is out of bounds #

错误根源 #

这个错误的核心问题是掩码中的无效类别索引。TGS盐识别是二分类任务（盐/非盐），正常目标掩码的像素值应该只有0（非盐）和1（盐），但你的掩码里存在值为255的像素——Fastai默认的损失函数（比如交叉熵）期望类别索引是从0开始的连续整数，255显然超出了二分类的有效范围（0和1）。

你设置的open_mask(div=True)只是将掩码值除以255转为0-1的浮点数，但如果原始掩码里有255，处理后会变成1.0，可如果后续没有明确把255映射到正确的类别，或者损失函数依然按类别索引逻辑处理，就会触发越界错误。另外num_workers=0只是解决了数据加载的多线程冲突，和这个类别越界问题无关。

修复步骤 #

1. 先确认掩码数据的异常值：用代码检查掩码中的像素值分布：

   import numpy as np
   from fastai.vision import open_mask
   
   sample_mask = open_mask("path/to/your/mask.png", div=True).data.numpy()
   print("掩码中的唯一值：", np.unique(sample_mask))
   

   如果输出里包含1.0（对应原始的255），就需要针对性处理。
2. 自定义掩码加载函数：修改open_mask的逻辑，把255对应的1.0映射为正确的类别（如果255代表盐区域，就保留1；如果是无效标注，可设为0或作为忽略标签）：

   from fastai.vision import ImageSegment, open_mask
   import torch
   
   def open_tgs_mask(fn):
       mask = open_mask(fn, div=True).data
       # 把原始255转来的1.0确认为盐类别（如果是无效区域则改为0）
       mask = torch.where(mask == 1.0, torch.tensor(1.0), torch.tensor(0.0))
       return ImageSegment(mask)
   

3. 明确数据块的类别设置：在构建数据时指定类别为[0,1]，确保损失函数适配二分类：

   data = (SegmentationItemList.from_df(df, path)
           .split_by_rand_pct()
           .label_from_func(get_mask_fn, classes=[0, 1])
           .transform(get_transforms(), tfm_y=True)
           .databunch(bs=16, num_workers=0)
           .normalize(imagenet_stats))
   

GPU环境错误：CUDA error: device-side assert triggered #

错误根源 #

这个GPU端的断言错误本质上是CPU错误的延续——GPU上的错误提示不会像CPU那样直接指出类别越界，而是以断言触发的形式表现。另外，你使用的自定义IOU指标如果存在逻辑漏洞（比如张量形状不匹配、除以零、未正确处理GPU张量），也会触发这个错误。

修复步骤 #

1. 优先解决CPU端的掩码类别问题：按照上面的CPU修复步骤处理完掩码数据，确保所有目标值都是0或1，这是解决GPU错误的前提。
2. 检查并修正自定义IOU指标：IOU计算时要注意几个关键点：
   • 保证预测和目标的形状完全匹配；
   • 避免除以零的情况（比如当目标和预测都没有正样本时）；
   • 全程用张量运算，不要随意在GPU和CPU之间切换张量。
     给你一个适配二分类的IOU指标示例：

   import torch
   
   def tgs_iou(input, target):
       # 把模型输出的通道维度转为类别索引
       input = torch.argmax(input, dim=1).float()
       target = target.squeeze(1).float()  # 确保目标和预测形状一致
       
       intersection = (input * target).sum(dim=[1,2])
       union = input.sum(dim=[1,2]) + target.sum(dim=[1,2]) - intersection
       
       # 处理union为0的情况（避免除以零）
       iou = torch.where(union == 0, torch.tensor(1.0).to(input.device), intersection / union)
       return iou.mean()  # 返回批次的平均IOU
   

   然后在构建模型时注册这个指标：

   learn = unet_learner(data, models.resnet34, metrics=[tgs_iou])
   

3. 验证数据加载的一致性：确保图像和掩码应用了相同的变换（tfm_y=True），避免出现形状不匹配的情况，这也是GPU断言错误的常见诱因。

额外提示 #

• Fastai 1.0.60是比较早期的版本，如果上述修复后仍有问题，可以考虑升级到Fastai 1.x的最新稳定版（注意部分API可能有细微调整）；
• 对于TGS这种二分类分割任务，也可以尝试用LabelList替代SegmentationItemList，直接把掩码作为二分类标签处理，可能会简化数据流程。

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