上传权重的 torch hub yolov5 模型检测尺寸不正确

Question

我正在尝试在本地使用 yolov5n6 预训练模型，但似乎存在一些错误。

我希望检测大小是 torch.Size(3,4)，因为我分批向模型发送了 3 次相同的图像。我在每个打印功能旁边打印了结果和复制粘贴。

获得意外输出的可能原因是什么？

import torch
import cv2

X = cv2.imread("/path/to/image/normalized_breast.PNG")
X = cv2.resize(X,(1024,1024))
 
print(X.shape) # (1024, 1024, 3)

model = torch.hub.load('ultralytics/yolov5', 'yolov5n6', pretrained=True, classes=1)
checkpoint_ = torch.load('/path/to/pretrained/model/yolov5_breast.pt')['model']
model.load_state_dict(checkpoint_.state_dict())

img = torch.from_numpy(X/255).permute(2,0,1).float().unsqueeze(0)
batch = torch.concatenate([img,img,img]).to("cuda")

print(batch.size()) # torch.Size([3, 3, 1024, 1024])

with torch.no_grad():
    model.eval()
    detections = model(batch)

print(len(detections),         # 2
      detections[0].size(),      # torch.Size([3, 65280, 6])
      detections[1][0].size()) # torch.Size([3, 3, 128, 128, 6])