【计算机视觉】YOLOv5训练报错：RuntimeError: result type Float can‘t be cast to the desired output type __int64

一、问题提出

在使用yolov5训练自己的数据集时，运行train.py时出现上述问题：

Image sizes 1536 train, 1536 val
Using 2 dataloader workers
Logging results to runs\train\exp6
Starting training for 100 epochs...

     Epoch   gpu_mem       box       obj       cls    labels  img_size
  0%|                                                                                                                                                                    | 0/3236 [00:00<?, ?it/s] 
Traceback (most recent call last):
  File "train.py", line 630, in <module>
    main(opt)
  File "train.py", line 527, in main
    train(opt.hyp, opt, device, callbacks)
  File "train.py", line 324, in train
    loss, loss_items = compute_loss(pred, targets.to(device))  # loss scaled by batch_size
  File "D:\CodeProject\My-Models\utils\loss.py", line 142, in __call__
    tcls, tbox, indices, anchors = self.build_targets(p, targets)  # targets
  File "D:\CodeProject\My-Models\utils\loss.py", line 240, in build_targets
    indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))  # image, anchor, grid indices
RuntimeError: result type Float can't be cast to the desired output type __int64

二、解决问题

将loss.py中：

gain = torch.ones(7, device=targets.device)

改为：

gain = torch.ones(7, device=targets.device).long()

原因是新版本的torch无法自动执行此转换，旧版本torch可以。

修改源码的位置为：

    def build_targets(self, p, targets):
        # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
        na, nt = self.na, targets.shape[0]  # number of anchors, targets
        tcls, tbox, indices, anch = [], [], [], []
        # gain = torch.ones(7, device=targets.device)  # normalized to gridspace gain
        gain = torch.ones(7, device=targets.device).long()
        ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt)  # same as .repeat_interleave(nt)
        targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2)  # append anchor indices