训练YOLOV7-w6模型笔记

RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

在实验室的服务器上运行下面指令，遇到上述问题

python train_aux.py --weights weights/yolov7-w6_training.pt --cfg cfg/training/yolov7-w6-Helmet.yaml --data data/Helmet.yaml --device 0 --batch-size 8 --epochs 100 --img-size 1280 1024

环境配置为：torch 1.13.1+cu116 CUDA:0 (NVIDIA GeForce RTX 3090,乌邦图系统）

因为用的是自己的数据集，遇到索引张量不在同一个设备（GPU/CPU）上的问题

从github上找到的解决办法，亲测有效，附原地址：Yolov7-w6.pt custom training runtime error indices should be either in cpu or on the same device · Issue #1228 · WongKinYiu/yolov7 · GitHub

主要修改以下两处地方，均在/utils/loss.py文件下

1.首先是loss.py代码1330行和1483行附近

#更改前
all_anch.append(anch[i][idx])
from_which_layer.append(torch.ones(size=(len(b),)) * i)

将其中from_which_layer.append(torch.ones(size=(len(b),)) * i)均替换为：

from_which_layer.append((torch.ones(size=(len(b),)) * i).to('cuda'))

#更改后
all_anch.append(anch[i][idx])
from_which_layer.append((torch.ones(size=(len(b),)) * i).to('cuda'))

2.然后是loss.py代码1401行和1555行附近

#更改前
fg_mask_inboxes = matching_matrix.sum(0) > 0.0
matched_gt_inds = matching_matrix[:, fg_mask_inboxes].argmax(0)

在这两行中均添加一行代码：fg_mask_inboxes = fg_mask_inboxes.to(torch.device('cuda'))

#更改后
fg_mask_inboxes = matching_matrix.sum(0) > 0.0
fg_mask_inboxes = fg_mask_inboxes.to(torch.device('cuda'))
matched_gt_inds = matching_matrix[:, fg_mask_inboxes].argmax(0)

跑yolov7 p5的模型没事，跑p6需要更改下loss.py，因此记录以下。