一文详解如何实现PyTorch模型编译

作者:神经星星 时间:2023-06-03 02:46:27 

准备

本篇文章译自英文文档 Compile PyTorch Models。

作者是 Alex Wong。

更多 TVM 中文文档可访问 →TVM 中文站。

本文介绍了如何用 Relay 部署 PyTorch 模型。

首先应安装 PyTorch。此外,还应安装 TorchVision,并将其作为模型合集 (model zoo)。

可通过 pip 快速安装:

pip install torch==1.7.0
pip install torchvision==0.8.1

或参考官网:pytorch.org/get-started…

PyTorch 版本应该和 TorchVision 版本兼容。

目前 TVM 支持 PyTorch 1.7 和 1.4,其他版本可能不稳定。

import tvm
from tvm import relay
import numpy as np
from tvm.contrib.download import download_testdata
# 导入 PyTorch
import torch
import torchvision

加载预训练的 PyTorch 模型

model_name = "resnet18"
model = getattr(torchvision.models, model_name)(pretrained=True)
model = model.eval()
# 通过追踪获取 TorchScripted 模型
input_shape = [1, 3, 224, 224]
input_data = torch.randn(input_shape)
scripted_model = torch.jit.trace(model, input_data).eval()
输出结果:

Downloading: "download.pytorch.org/models/resn…" to /workspace/.cache/torch/hub/checkpoints/resnet18-f37072fd.pth

0%| | 0.00/44.7M [00:00<?, ?B/s] 11%|# | 4.87M/44.7M [00:00<00:00, 51.0MB/s] 22%|##1 | 9.73M/44.7M [00:00<00:00, 49.2MB/s] 74%|#######3 | 32.9M/44.7M [00:00<00:00, 136MB/s] 100%|##########| 44.7M/44.7M [00:00<00:00, 129MB/s]

加载测试图像

经典的猫咪示例:

from PIL import Image
img_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true"
img_path = download_testdata(img_url, "cat.png", module="data")
img = Image.open(img_path).resize((224, 224))
# 预处理图像,并将其转换为张量
from torchvision import transforms
my_preprocess = transforms.Compose(
[
       transforms.Resize(256),
       transforms.CenterCrop(224),
       transforms.ToTensor(),
       transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
)
img = my_preprocess(img)
img = np.expand_dims(img, 0)

将计算图导入 Relay

将 PyTorch 计算图转换为 Relay 计算图。input_name 可以是任意值。

input_name = "input0"
shape_list = [(input_name, img.shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)

Relay 构建

用给定的输入规范,将计算图编译为 llvm target。

target = tvm.target.Target("llvm", host="llvm")
dev = tvm.cpu(0)
with tvm.transform.PassContext(opt_level=3):
   lib = relay.build(mod, target=target, params=params)

输出结果:

/workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
 "target_host parameter is going to be deprecated. "

在 TVM 上执行可移植计算图

将编译好的模型部署到 target 上:

from tvm.contrib import graph_executor
dtype = "float32"
m = graph_executor.GraphModule(lib["default"](dev))
# 设置输入
m.set_input(input_name, tvm.nd.array(img.astype(dtype)))
# 执行
m.run()
# 得到输出
tvm_output = m.get_output(0)

查找分类集名称

在 1000 个类的分类集中,查找分数最高的第一个:

synset_url = "".join(
[
"https://raw.githubusercontent.com/Cadene/",
"pretrained-models.pytorch/master/data/",
"imagenet_synsets.txt",
]
)
synset_name = "imagenet_synsets.txt"
synset_path = download_testdata(synset_url, synset_name, module="data")
with open(synset_path) as f:
   synsets = f.readlines()
synsets = [x.strip() for x in synsets]
splits = [line.split(" ") for line in synsets]
key_to_classname = {spl[0]: " ".join(spl[1:]) for spl in splits}
class_url = "".join(
[
"https://raw.githubusercontent.com/Cadene/",
"pretrained-models.pytorch/master/data/",
"imagenet_classes.txt",
]
)
class_name = "imagenet_classes.txt"
class_path = download_testdata(class_url, class_name, module="data")
with open(class_path) as f:
   class_id_to_key = f.readlines()
class_id_to_key = [x.strip() for x in class_id_to_key]
# 获得 TVM 的前 1 个结果
top1_tvm = np.argmax(tvm_output.numpy()[0])
tvm_class_key = class_id_to_key[top1_tvm]
# 将输入转换为 PyTorch 变量,并获取 PyTorch 结果进行比较
with torch.no_grad():
   torch_img = torch.from_numpy(img)
   output = model(torch_img)
# 获得 PyTorch 的前 1 个结果
   top1_torch = np.argmax(output.numpy())
   torch_class_key = class_id_to_key[top1_torch]
print("Relay top-1 id: {}, class name: {}".format(top1_tvm, key_to_classname[tvm_class_key]))
print("Torch top-1 id: {}, class name: {}".format(top1_torch, key_to_classname[torch_class_key]))

输出结果:

Relay top-1 id: 281, class name: tabby, tabby cat
Torch top-1 id: 281, class name: tabby, tabby cat

下载 Python 源代码:from_pytorch.py

下载 Jupyter Notebook:from_pytorch.ipynb

来源:https://juejin.cn/post/7225898744449155130

标签:PyTorch,模型,编译
0
投稿

猜你喜欢

  • Python设计模式中的创建型工厂模式

    2023-05-09 17:47:02
  • Python爬虫框架Scrapy基本用法入门教程

    2021-08-17 19:50:45
  • Cpython3.9源码解析python中的大小整数

    2021-11-17 05:45:42
  • ASP制作中使用MYSQL的分析

    2008-10-13 09:25:00
  • js获取select选中的option的text示例代码

    2024-04-19 09:58:48
  • Python调用百度OCR实现图片文字识别的示例代码

    2021-04-07 14:21:50
  • python3 tkinter实现点击一个按钮跳出另一个窗口的方法

    2023-10-27 01:19:11
  • 简单探讨一下python线程锁

    2021-03-24 12:06:54
  • python 猴子补丁(monkey patch)

    2022-12-25 02:15:13
  • lhgcalendar时间插件限制只能选择三个月的实现方法

    2024-04-10 13:56:53
  • 对Python中DataFrame按照行遍历的方法

    2023-01-03 23:08:59
  • asp如何远程读取数据库页面?

    2010-06-16 09:57:00
  • django中ImageField的使用详解

    2023-09-28 03:58:37
  • Python绘制牛奶冻曲线(高木曲线)案例

    2022-04-05 16:35:08
  • asp怎么实现中文字符串按声母检索

    2010-05-16 21:19:00
  • python发送伪造的arp请求

    2022-11-24 00:47:35
  • Python对XML文件实现增删改查操作

    2023-11-19 20:42:03
  • Python中装饰器高级用法详解

    2023-09-30 09:46:42
  • MYSQL自定义函数判断是否正整数的实例代码

    2024-01-19 19:14:20
  • python用pickle模块实现“增删改查”的简易功能

    2022-01-13 06:38:56
  • asp之家 网络编程 m.aspxhome.com