如何在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 (most recent call last) in 3 import torch 4 ----> 5 model = torch.load(os.getcwd()+'/weights/last.pt')

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

Answer 1

为了加载模型的权重，您应该首先导入模型脚本。我猜它位于 /weights/last.py。之后，您可以加载模型的权重。

示例代码可能如下：

import os  
import torch 
from weights.last import Model # I assume you named your model as Model, change it accordingly

model = Model()  # Then in here instantiate your model
model.load_state_dict(torch.load(ospath.join(os.getcwd()+'/weights/last.pt')))  # Then load your model's weights.

 

# 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)

在此解决方案中，不要忘记您应该运行您当前工作目录中的程序，如果您运行来自权重文件夹的程序，您可能会收到错误。

Answer 2

您应该可以在这个目录中找到权重：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

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

How to load custom model in pytorch

python

pytorch

yolo

yolov5