torch.hub.load()函数的使用——联网加载权重以及如何加载本地权重

torch.hub.load()函数可以通过两种方式加载模型。

pytorch官方文档

函数原型：

torch.hub.load(repo_or_dir, model, *args, source='github', force_reload=False, verbose=True, skip_validation=False, **kwargs)

原函数

参数说明：

repo_or_dir ( string ) – 如果source是 ‘github’，这应该对应于repo_owner/repo_name[:tag_name]具有可选标签/分支格式的 github 存储库，例如 ‘pytorch/vision:0.10’。如果tag_name未指定，则假定默认分支为main存在，否则为master。如果source是“local”，则它应该是本地目录的路径。

model ( string ) – 在 repo/dir’s 中定义的可调用（入口点）的名称hubconf.py。

*args（可选）– callable 的相应参数model。

source ( string , optional ) – ‘github’ 或 ‘local’。指定如何 repo_or_dir解释。默认为“github”。

force_reload ( bool , optional ) – 是否无条件强制重新下载github repo。如果 没有任何影响 source = ‘local’。默认为False

verbose ( bool , optional ) – 如果False，静音有关命中本地缓存的消息。请注意，有关首次下载的消息无法静音。如果source = 'local’没有任何影响。默认为True。

skip_validation ( bool , optional ) – 如果False，torchhub 将检查github参数指定的分支或提交是否正确属于 repo 所有者。这将向 GitHub API 发出请求；您可以通过设置GITHUB_TOKEN环境变量来指定非默认 GitHub 令牌 。默认为False。

**kwargs (可选) – callable 的相应 kwargs model。

实践干货：

1、联网

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

就可以直接从网上下载模型了

2、加载本地模型

2.1 文件下载

注意torch.hub.load()它实现的是加载权重文件，要预测还需要模型等等文件，model = torch.hub.load('ultralytics/yolov5', 'yolov5s')中的'ultralytics/yolov5'即是除了权重外的其它文件，方法1可以直接拉取，那么本地加载则也需要自行下载。

同样以YOLO5为例。

首先下载官方的YOLO5源码

除了YOLO5源码（下载后不需要训练），还需要下载权重文件。

2.2 参数赋值

如果YOLO5源码的存储地址如图所示：

则函数的第一个参数repo_or_dir 为G:\jinx\yolov5，

如果权重文件目录如图所示：

权重地址参数为G:\jinx\yolov5\pt\yolo5s.pt

2.3 总和

model = torch.hub.load('G:\jinx\yolov5', 'custom', 'G:\jinx\yolov5\pt\yolo5s.pt', source='local')

3、可能bug

问题1："Arial.ttf"无法下载

在yolov5/utils/plots.py中，Annotator类的check_font()函数注释掉即可。

本网站所收集的部分公开资料来源于AI生成和互联网，转载的目的在于传递更多信息及用于网络分享，并不代表本站赞同其观点和对其真实性负责，也不构成任何其他建议。