一、问题提出

在使用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
GitHub 加速计划 / yo / yolov5
510
37
下载
yolov5 - Ultralytics YOLOv8的前身,是一个用于目标检测、图像分割和图像分类任务的先进模型。
最近提交(Master分支:2 天前 )
e62a31b6 Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> 13 天前
882c35fc Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> 26 天前
Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐