如何在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(最近调用最后(3进口火炬4.---->5型号=火炬加载(os.getcwd((+'/wights/last.pt'(

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

Answer 1

您应该能够在以下目录中找到权重：yolov5/runs/strain/exp/wights/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

Answer 2

为了加载模型的权重，您应该首先导入模型脚本。我猜它位于/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)

在这个解决方案中，不要忘记您应该从当前工作目录运行程序，如果您从weights文件夹运行程序，您可能会收到错误。

Answer 3

如果你想加载本地保存的模型，你可以尝试这个

import torch
model = torch.hub.load('.', 'custom', 'yourmodel.pt', source='local')