Torch

hub

2023-04-23

[[torch]] 中用于加载模型的函数，可以加载 [[github]] 仓库中的模型，也可以加载本地模型。

方法

list

查看可用模型

1
import torch
2


3
torch.hub.list("ultralytics/yolov5")
4
# 返回: ['custom', 'yolov5l', 'yolov5l6', 'yolov5m', 'yolov5m6', 'yolov5n', 'yolov5n6', 'yolov5s', 'yolov5s6', 'yolov5x', 'yolov5x6']

help

查看模型帮助

1
import torch
2


3
torch.hub.help("pytorch/vision", "resnet18")

load

加载模型

1
import torch
2


3
# 加载
4
# github 仓库和模型
5
model = torch.hub.load("ultralytics/yolov5", "yolov5s")
6


7
# 本地模型
8
model = torch.hub.load(
9
                       "ultralytics/yolov5",
10
                       "custom",
11
                       path="path/best.pt"
12
                       )
13


14
# 本地版本和模型
15
model = torch.hub.load(
16
                       "path/yolov5",
17
                       "custom",
18
                       path="path/best.pt",
19
                       source="local"
20
                       )

参数

  • device: GPU
  • _verbose: 静默加载

设置

1
# 设置
2
model.conf = 0.25   # 置信度
3
model.classes = [0] # 分类

使用

1
res = model("file.jpg")
2


3
res.print()
4
res.save()
5
res.show()  # 预览图片
6
res.names: 所有 classess
7
res.xyxy[0]
8
res.files: 文件名 # 列表形式
9


10
# pandas 输出
11
res.pandas().xyxy[0]
12
res.pandas().xyxy[0].sort_values("xmin") # 排序从左到右
13
# json 输出
14
res.pandas().xyxy[0].to_json(orient="records")

pandas 包含:

  • names: 所有 classess #字典
  • files: 文件名 #列表
  • xyxy
  • xyxyn: 包含归一化
  • xywhn: 归一化
  • xywhn
  • n
  • t
  • s

加载到设备

1
model.cpu()  # CPU
2
model.cuda()  # GPU
3
model.to(device)  # i.e. device=torch.device(0)

download_url_to_file

1
import torch
2


3
# 下载文件到本地
4
torch.hub.download_url_to_file('http://url/img.jpg', '/save_path/file_name')

参考

# Python
# Torch
# YOLO