如何在pytorch中加载定制模型

7rtdyuoh  于 2022-11-23  发布在  其他
关注(0)|答案(3)|浏览(155)

我尝试加载我的预训练模型(yolov 5 n)并在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)

我得到了以下错误:
模块未找到错误跟踪(最近调用最后)在3导入 Torch 4 ----〉5模型= Torch .load(os.getcwd()+'/weights/last.pt')
我的模型位于文件夹/weights/last.py中,我不确定我做错了什么。请问,我的代码中缺少了什么。

b91juud3

b91juud31#

您应该可以在以下目录中找到权重:yolov 5/跑步/训练/经验/重量/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后从该目录加载自定义模型

yshpjwxd

yshpjwxd2#

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

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

vc6uscn9

vc6uscn93#

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

import torch

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

相关问题