使用F.cross_entropy触发CUDA设备端断言错误的排查求助

问题复现代码 #

import torch
device = "cuda" if torch.cuda.is_available() else "cpu"

a = torch.Tensor([[-10.3353, -28.4371,   2.0768,   -4.2789,  -8.6644,  -6.0815],
        [-10.3353, -28.4371,   2.0768,   -4.2789,  -8.6644,  -6.0815],
        [-10.3353, -28.4371,   2.0768,   -4.2789,  -8.6644,  -6.0815],
        [-10.3353, -28.4371,   2.0768,   -4.2789,  -8.6644,  -6.0815]]).to(device)
b = torch.Tensor([ -100,  -1,  -100,  2456]).long().to(device)

loss = torch.nn.functional.cross_entropy(a,b,ignore_index=-1)
print(loss)

报错信息 #

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-9-29fe2d43a573> in <cell line: 1>()
----> 1 loss = torch.nn.functional.cross_entropy(a,b,ignore_index=-1).to(device)
      2 loss

/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
   3027     if size_average is not None or reduce is not None:
   3028         reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 3029     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
   3030 
   3031 

RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

错误原因 #

PyTorch的cross_entropy函数对target张量的取值有严格要求：

1. 有效类别索引必须在0到num_classes-1范围内（这里input张量a的最后一维维度是6，所以有效索引是0-5）；
2. 需要忽略的样本索引必须等于ignore_index指定的值（此处为-1）。

你的target张量b中存在两个无效值：

• -100不等于设置的ignore_index=-1，会被当作类别索引处理，但超出有效范围；
• 2456远大于num_classes-1=5，直接触发索引越界的设备端断言。

解决方法 #

修正target张量的取值，确保所有值要么是0-5之间的有效类别索引，要么等于ignore_index=-1。

修正后的代码示例

import torch
device = "cuda" if torch.cuda.is_available() else "cpu"

a = torch.Tensor([[-10.3353, -28.4371,   2.0768,   -4.2789,  -8.6644,  -6.0815],
        [-10.3353, -28.4371,   2.0768,   -4.2789,  -8.6644,  -6.0815],
        [-10.3353, -28.4371,   2.0768,   -4.2789,  -8.6644,  -6.0815],
        [-10.3353, -28.4371,   2.0768,   -4.2789,  -8.6644,  -6.0815]]).to(device)
# 将无效的-100替换为ignore_index=-1，2456修正为有效索引3
b = torch.Tensor([ -1,  -1,  -1,  3]).long().to(device)

loss = torch.nn.functional.cross_entropy(a,b,ignore_index=-1)
print(loss)

其他可选方案

• 如果-100也是需要忽略的样本，可将ignore_index设置为-100，同时修正2456为有效索引；
• 检查数据预处理流程，确保生成target时不会产生超出类别范围的无效值。

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