训练LayoutLM触发CUDA设备断言错误的原因排查

我正在使用自定义收据数据集训练LayoutLM模型，在开启GPU P100的Kaggle环境中运行代码时，遭遇了CUDA设备端断言触发的错误。该代码在仅包含4个标签的SROIE数据集上可正常运行，但切换到含8个标签的自定义数据集时出现如下报错：

/opt/conda/conda-bld/pytorch_1603729138878/work/aten/src/ATen/native/cuda/Indexing.cu:658: indexSelectLargeIndex: block: [160,0,0], thread: [125,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
/opt/conda/conda-bld/pytorch_1603729138878/work/aten/src/ATen/native/cuda/Indexing.cu:658: indexSelectLargeIndex: block: [160,0,0], thread: [126,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
/opt/conda/conda-bld/pytorch_1603729138878/work/aten/src/ATen/native/cuda/Indexing.cu:658: indexSelectLargeIndex: block: [160,0,0], thread: [127,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
Iteration:   2%|▋                               | 3/140 [00:01<01:25,  1.61it/s]
Epoch:   0%|                                              | 0/1 [00:01<?, ?it/s]
Traceback (most recent call last):
  File "run_seq_labeling.py", line 832, in <module>
    main()
  File "run_seq_labeling.py", line 725, in main
    args, train_dataset, model, tokenizer, labels, pad_token_label_id
  File "run_seq_labeling.py", line 240, in train
    outputs = model(**inputs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/layoutlm/modeling/layoutlm.py", line 224, in forward
    head_mask=head_mask,
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/layoutlm/modeling/layoutlm.py", line 178, in forward
    input_ids, bbox, position_ids=position_ids, token_type_ids=token_type_ids
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/layoutlm/modeling/layoutlm.py", line 102, in forward
    + token_type_embeddings
RuntimeError: CUDA error: device-side assert triggered

错误触发原因 #

• 模型分类头维度不匹配：切换到8标签数据集后，未重新初始化或更新LayoutLM的分类头，模型仍保留SROIE数据集的4分类输出维度。当自定义数据集中的标签索引（最大为7）尝试匹配模型的4维输出时，触发索引越界断言。
• 标签映射未同步更新：自定义数据集的标签转索引逻辑生成的索引值（0-7）超出了模型当前支持的类别范围（0-3），导致CUDA在执行索引选择操作时，出现srcIndex >= srcSelectDimSize的违规情况。
• 预处理参数遗漏更新：pad_token_label_id或其他标签相关预处理参数未随新数据集调整，部分样本的标签编码出现超出模型输出维度的数值，进而触发设备端断言。

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