CUDA Tensor 越界错误的分析与解决方案 || SoftDiceLoss || y_onehot.scatter_(1, gt, 1)

在运行深度学习训练过程中，以下错误表明发生了 CUDA 张量索引越界 的问题：

../aten/src/ATen/native/cuda/ScatterGatherKernel.cu:367: operator(): block: [1227,0,0], thread: [90,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/ScatterGatherKernel.cu:367: operator(): block: [1227,0,0], thread: [91,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/ScatterGatherKernel.cu:367: operator(): block: [1227,0,0], thread: [92,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/ScatterGatherKernel.cu:367: operator(): block: [1227,0,0], thread: [93,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/ScatterGatherKernel.cu:367: operator(): block: [1227,0,0], thread: [94,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && "index out of bounds"` failed.
  0%|          | 0/250 [00:01<?, ?it/s]
Traceback (most recent call last):
  File "/media/sunjiakuo/78de387e-1660-4914-b70a-224cd2c430e9/xtx/teacher/train_unet.py", line 82, in <module>
    train_loss, train_dice = trainer.train_epoch()
  File "/media/sunjiakuo/78de387e-1660-4914-b70a-224cd2c430e9/xtx/teacher/SegTrainer.py", line 173, in train_epoch
    loss = self.loss_seg(seg_pred_batch, mask_batch)
  File "/home/sunjiakuo/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/media/sunjiakuo/78de387e-1660-4914-b70a-224cd2c430e9/xtx/teacher/MyLoss.py", line 149, in forward
    tp, fp, fn, _ = get_tp_fp_fn_tn(net_output, gt, axes, loss_mask, False)
  File "/media/sunjiakuo/78de387e-1660-4914-b70a-224cd2c430e9/xtx/teacher/MyLoss.py", line 98, in get_tp_fp_fn_tn
    y_onehot.scatter_(1, gt, 1)
RuntimeError: CUDA error: device-side assert triggered

这种错误通常由 张量的索引值超出定义范围 引起。在代码中，错误栈指向了以下问题点：

1. 错误代码定位：

   • 报错的主要函数是 y_onehot.scatter_(1, gt, 1)，其中 gt（ground truth）中的值超出了张量的索引范围。
   • scatter_ 函数在索引时需要确保 gt 的值位于预定义范围内，否则会触发 “index out of bounds” 错误。

2. 问题起因：

   • 数据集中 标签的保存格式 是 RGB（即像素值范围为 [0, 255]），而非常规的二值或多分类标签格式（例如 [0, 1] 或 [0, num_classes - 1]）。
   • 因此，在使用 scatter_ 函数时，gt 的索引值为 255（或其他超出范围的值），导致张量操作失败。

---

解决方案

为了解决这个问题，可以通过以下步骤将 gt 的值限制在有效索引范围内：

1. 使用 clamp 函数：

   • 将标签的值限制在 [0, 1] 范围内：

   gt = gt.clamp(0, 1)
   

   • 该函数的作用是将 gt 中的所有值 小于 0 的值设置为 0，超过 1 的值设置为 1。对于 RGB 数据中的标签值（255），会被转化为 1，从而避免索引越界。

2. 修正代码：

   • 在调用 y_onehot.scatter_ 之前，插入以下代码来规范化 gt 的值：

   gt = gt.clamp(0, 1)  # 将标签值限制在 0-1 区间
   

3. 完整代码示例：

   def get_tp_fp_fn_tn(net_output, gt, axes, loss_mask, some_flag):
       # 在 scatter_ 操作之前，确保 gt 的索引值在合法范围内
       gt = gt.clamp(0, 1)
       y_onehot.scatter_(1, gt, 1)
       # 其他代码逻辑...
   

---

背景说明

• 问题来源：

  • 您提到在 重新建立 nii 数据集时，使用了 RGB 格式保存标签，导致标签值范围是 [0, 255]。
  • 常规深度学习标签值通常是整数类型，并且范围为 [0, num_classes - 1]。这种保存方式与现有的模型输入格式不匹配，导致了索引越界问题。

• 修正目标：

  • 将 RGB 标签值进行 归一化处理 或 直接映射为二值/多分类标签格式。
  • clamp 的作用是快速限制值的范围，避免错误，但更优雅的方式是将数据集转换为标准格式。

---

参考文献与补充学习

关于 SoftDiceLoss，可以参考以下资源：

• SoftDiceLoss 代码实现
• SoftDiceLoss 的核心思想是通过计算预测结果与真实标签的交集，来优化分割模型的表现。它尤其适用于类别不平衡的分割任务。