训练Faster R-CNN模型时出现KeyError: tensor(1)错误求助

训练命令 #

!python train.py --data data_configs/custom_data.yaml --epochs 1 --imgsz 850 --model fasterrcnn_resnet50_fpn_v2 --name custom_training --batch 4 --device cuda

报错信息（中文翻译） #

File "C:\Users\..\fastercnn-pytorch-training\train.py", line 574, in <module>
    main(args)

  File "C:\Users\..\fastercnn-pytorch-training\fastercnn-pytorch-training\train.py", line 414, in main
    batch_loss_rpn_list = train_one_epoch(

  File "C:\Users\..\fastercnn-pytorch-training\fastercnn-pytorch-training\torch_utils\engine.py", line 46, in train_one_epoch
    for images, targets in metric_logger.log_every(data_loader, print_freq, header):
  
File "C:\Users\..\fastercnn-pytorch-training\fastercnn-pytorch-training\torch_utils\utils.py", line 174, in log_every
    for batch in iterable:
 
 File "C:\Users\..\AppData\Roaming\Python\Python310\site-packages\torch\utils\data\dataloader.py", line 701, in __next__
    data = self._next_data()
 
 File "C:\Users\..\AppData\Roaming\Python\Python310\site-packages\torch\utils\data\dataloader.py", line 757, in _next_data
    data = self._dataset_fetcher.fetch(index)  # 可能触发StopIteration
 
 File "C:\Users\..\AppData\Roaming\Python\Python310\site-packages\torch\utils\data\_utils\fetch.py", line 52, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
 
 File "C:\Users\..\AppData\Roaming\Python\Python310\site-packages\torch\utils\data\_utils\fetch.py", line 52, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
 
 File "C:\Users\..\fastercnn-pytorch-training\fastercnn-pytorch-training\datasets.py", line 317, in __getitem__
    sample = self.transforms(image=image_resized,
  
File "C:\Users\..\AppData\Roaming\Python\Python310\site-packages\albumentations\core\composition.py", line 421, in __call__
    self.preprocess(data)
  
File "C:\Users\..\AppData\Roaming\Python\Python310\site-packages\albumentations\core\composition.py", line 447, in preprocess
    p.preprocess(data)
 
 File "C:\Users\..\AppData\Roaming\Python\Python310\site-packages\albumentations\core\utils.py", line 171, in preprocess
    data = self.add_label_fields_to_data(data)
  
File "C:\Users\..\AppData\Roaming\Python\Python310\site-packages\albumentations\core\utils.py", line 220, in add_label_fields_to_data
    data[data_name] = self._process_label_fields(data, data_name)
  
File "C:\Users\..\AppData\Roaming\Python\Python310\site-packages\albumentations\core\utils.py", line 229, in _process_label_fields
    encoded_labels = self._encode_label_field(data, data_name, label_field)
 
 File "C:\Users\..\AppData\Roaming\Python\Python310\site-packages\albumentations\core\utils.py", line 249, in _encode_label_field
    encoded_labels = encoder.fit_transform(data[label_field]).reshape(-1, 1)
 
 File "C:\Users\...\AppData\Roaming\Python\Python310\site-packages\albumentations\core\utils.py", line 84, in fit_transform
    return self.transform(y)
 
 File "C:\Users\...\AppData\Roaming\Python\Python310\site-packages\albumentations\core\utils.py", line 80, in transform
    return np.array([self.classes_[label] for label in y])
  
File "C:\Users\...\AppData\Roaming\Python\Python310\site-packages\albumentations\core\utils.py", line 80, in <listcomp>
    return np.array([self.classes_[label] for label in y])

KeyError: tensor(1)

错误原因 #

核心问题是Albumentations标签编码器无法识别tensor类型的标签值1：

1. 编码器内部用字典映射类别，键为Python原生int类型，但传入的标签是PyTorch tensor对象，导致字典查找失败；
2. 或者数据集标注的类别ID在配置文件的类别列表中不存在。

排查步骤 #

• 检查custom_data.yaml配置：确认classes列表对应的类别ID从0开始，标注文件里的所有类别ID都在配置的类别范围内（比如yaml只定义1个类别时，ID不能是1）；
• 修正标签数据类型：在datasets.py的__getitem__方法（约317行附近），将传给transforms的标签从tensor转为Python整数，比如用.item()处理target['labels']；
• 验证Albumentations变换配置：如果用了ClassLabel类变换，确保classes参数和数据集类别列表完全匹配；
• 检查标注文件：确认所有标注的目标类别ID合法，无拼写错误或超出范围的ID。

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