YOLOv7训练自定义数据集出现RuntimeError设备不匹配问题求助

问题场景 #

已下载YOLOv7自定义数据集并克隆仓库，使用以下命令启动训练：

python train.py --workers 0 --batch-size 4 --device 0 --data data\face_mask.yaml --img 640 640 --cfg cfg\training\yolov7-face_mask.yaml --weights yolov7_training.pt --name yolov7-face-mask --hyp data\hyp.scratch.custom.yaml --epochs 300

初始运行正常，输出环境及参数信息后触发RuntimeError。

报错信息 #

Epoch   gpu_mem       box       obj       cls     total    labels  img_size
  0%|                                                                          | 0/160 [00:08<?, ?it/s]
Traceback (most recent call last):
  File "C:\YOLOv7\yolov7-gpu\train.py", line 616, in <module>
    train(hyp, opt, device, tb_writer)
  File "C:\YOLOv7\yolov7-gpu\train.py", line 363, in train
    loss, loss_items = compute_loss_ota(pred, targets.to(device), imgs)  # loss scaled by batch_size
  File "C:\YOLOv7\yolov7-gpu\utils\loss.py", line 585, in __call__
    bs, as_, gjs, gis, targets, anchors = self.build_targets(p, targets, imgs)
  File "C:\YOLOv7\yolov7-gpu\utils\loss.py", line 759, in build_targets
    from_which_layer = from_which_layer[fg_mask_inboxes]
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

软硬件环境 #

• 硬件：Acer Nitro 5笔记本、Intel i5-11处理器、NVIDIA RTX 3050 GPU
• 软件：Python 3.10、Anaconda、NVIDIA-SMI 517.48、驱动版本517.48、CUDA 11.7

错误原因 #

该错误源于张量设备不匹配：用于索引的fg_mask_inboxes张量与被索引的from_which_layer张量分别处于不同设备（一个在GPU，一个在CPU），导致PyTorch的索引操作无法执行。此问题多出现于YOLOv7的OTA损失计算逻辑中，部分张量未被统一到训练指定的设备上。

解决方案 #

方案1：修复loss.py中的设备同步 #

打开utils/loss.py，定位到build_targets函数内报错的第759行附近，在执行索引操作前强制将两个张量统一到同一设备：

# 在from_which_layer = from_which_layer[fg_mask_inboxes]前添加
fg_mask_inboxes = fg_mask_inboxes.to(from_which_layer.device)

保存修改后重新运行训练命令即可。

方案2：降级PyTorch版本 #

高版本PyTorch（如2.x系列）与YOLOv7旧版本存在设备兼容性问题，可降级到适配CUDA 11.7的PyTorch 1.13.x版本：

conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.7 -c pytorch -c nvidia

方案3：临时用CPU训练排查问题 #

若仅需验证数据集和配置文件的正确性，可修改训练命令指定使用CPU：

python train.py --workers 0 --batch-size 4 --device cpu --data data\face_mask.yaml --img 640 640 --cfg cfg\training\yolov7-face_mask.yaml --weights yolov7_training.pt --name yolov7-face-mask --hyp data\hyp.scratch.custom.yaml --epochs 300

注：CPU训练速度极慢，仅用于问题排查。

内容的提问来源于stack exchange，提问作者Tri Amin Ridho