如何在pytorch中加载自定义模型

Question

我正在尝试加载我的预训练模型 (yolov5n) 并在 PyTorch 中使用以下代码对其进行测试：

import os  
import torch 
model = torch.load(os.getcwd()+'/weights/last.pt')
 

# Images
imgs = ['https://example.com/img.jpg']   
# Inference
results = model(imgs)

# Results
results.print()
results.save()  # or .show()

results.xyxy[0]  # img1 predictions (tensor)
results.pandas().xyxy[0]  # img1 predictions (pandas)

我收到以下错误：

ModuleNotFoundError Traceback (最近一次调用最后) in 3 import torch 4 ----> 5 model = torch.load(os.getcwd()+'/weights/last.pt')

我的模型位于文件夹中/weights/last.py，我不确定我在做什么错误。你能告诉我，我的代码中缺少什么吗？

Answer 1

您应该能够在此目录中找到权重：yolov5/runs/train/exp/weights/last.pt

然后使用如下行加载权重：

model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5/runs/train/exp/weights/last.pt', force_reload=True) 

我有一个笔记本的示例，该笔记本在训练模型后从该目录加载自定义模型https://github.com/pylabel-project/samples/blob/main/pylabeler.ipynb