【PyTorch】存储YOLOv5的预训练模型 / torch.hub详解

存储YOLOv5的预训练模型 / torch.hub详解

下载源：https://github.com/ultralytics/yolov5/releases

构成：

• yolov5-6.1.zip，将文件解压在此文件夹下
• 后缀 .pt 文件
  • yolov5的与训练模型，目前本地可调用模型：
    [ ‘yolov5l.pt’, ‘yolov5l6.pt’, ‘yolov5m.pt’, ‘yolov5m6.pt’, ‘yolov5n.pt’, ‘yolov5n6.pt’, ‘yolov5s.pt’, ‘yolov5s6.pt’, ‘yolov5x.pt’, ‘yolov5x6.pt’]
  • 【之后有余力的话，可以对比几种不同YOLOv5预训练模型对无人机行人检测识别的影响 】

.pt文件

• .pt .pth .pkl 三种后缀文件只是后缀不同，格式一致，官方喜欢用.pt
• pytorch 模型文件，用以保存模型

---

./pytorch/vision 文件夹用途

• 存储PyTorch Hub发布的模型：torch.hub，
  • 一个包含计算机视觉、自然语言处理领域的诸多经典模型的聚合中心，让调用更方便
  • 为了调用各种经典机器学习模型，今后你不必重复造轮子了
  • 只需输入一行代码，就能实现一键调用。torch.hub.load()

---

PyTorch Hub允许用户对已发布的模型进行直接调用，无需自己训练模型：

参考来源：
https://blog.csdn.net/qq_38410428/article/details/101062287
https://blog.csdn.net/qq_34317565/article/details/122299172

1、查询可用的模型;
torch.hub.list()这个API列出repo中所有可用的入口点
【但是国内一般连接不上，报错如下，故而将模型文件放到本地。】
urllib.error.URLError: 由于连接方在一段时间后没有正确答复或连接的主机没有反应，连接尝试失败。>

2、加载模型;
调用上一步中查询到的计算机视觉模型，无需安装，直接进行如下调用即可，模型会自动下载到本地，无需手动下载。

  model = torch.hub.load('pytorch/vision', 'deeplabv3_resnet101', pretrained=True)
  自动进行【Downloading: "https://github.com/pytorch/vision/archive/main.zip" to C:\Users\流年/.cache\torch\hub\main.zip】      

【国内一般将模型放到本地，进行本地调用加载】
  path = '../../../../pytorch/vision/' #本地预训练模型路径
  model = torch.hub.load(path, 'yolov5s', source='local', pretrained=True)  # pretrained=True 设为预训练模式， source默认为github，因此需指定本地
  img = '../tmp/0000009.jpg'
  results = model(img) 	#用模型检测img图片，或者video，具体看模型的输入参数
  results.print() 		#输出检测结果 

 【torch.hub.load()参数具体调用信息直接查文档，下面列出几个重点 】
   torch.hub.load(repo_or_dir, model, *args, source='github', force_reload=False, verbose=True, skip_validation=False, **kwargs)
   - repo_or_dir：记得设置本地./pytorch/vision 路径
   - model 调用入口在：本地./pytorch/vision 路径下的hubconf.py文件里，因此模型名字要在hubconf.py中存在才能调用
   - source='local', 记得设置本地，默认：github
   - pretrained=True #设定预训练模式，默认False,为了代码清晰，最好还是加上参数赋值.
为了加快学习进度，训练的初期直接加载pretrain模型中预先训练好的参数
pretrained=True 加载网络结构和预训练参数
pretrained=False 只加载网络结构，不加载预训练参数，即不需要用预训练模型的参数来初始化。

3、查询模型中可用的方法。
从PyTorch Hub加载模型后，你可以用dir(model)查看模型的所有可用方法

4、可以自定义模型：

model = torch.hub.load('ultralytics/yolov5', 'custom', path='path/to/best.pt')  # local model
model = torch.hub.load('path/to/yolov5', 'custom', path='path/to/best.pt', source='local')  # local repo

---

保存模型与调用的几种方式：

来源：https://blog.csdn.net/weixin_38145317/article/details/103582549
方式一：
保存：

torch.save(model.state_dict(), mymodel.pth)#只保存模型权重参数，不保存模型结构
调用：
model = My_model(*args, **kwargs)  #这里需要重新模型结构，My_model
model.load_state_dict(torch.load(mymodel.pth)) #这里根据模型结构，调用存储的模型参数
model.eval()

方式二：
保存：

torch.save(model, mymodel.pth)#保存整个model的状态
调用：
model=torch.load(mymodel.pth)#这里已经不需要重构模型结构了，直接load就可以
model.eval()

---

Pytorch模型加载问题：
参考来源：
https://blog.csdn.net/qq_34097715/article/details/83345189
https://cloud.tencent.com/developer/article/1624133 Pytorch预训练模型以及修改

Pytorch模型的加载方式支持两种：一种是使用Pytorch自带的模型，另外一种是程序员自己编写的模型。