一文詳解如何實(shí)現(xiàn)PyTorch模型編譯

更新時間：2023年04月26日 09:09:29 作者：神經(jīng)星星

這篇文章主要為大家介紹了如何實(shí)現(xiàn)PyTorch?模型編譯詳解，有需要的朋友可以借鑒參考下，希望能夠有所幫助，祝大家多多進(jìn)步，早日升職加薪

準(zhǔn)備

本篇文章譯自英文文檔 Compile PyTorch Models。

作者是 Alex Wong。

更多 TVM 中文文檔可訪問 →TVM 中文站。

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

首先應(yīng)安裝 PyTorch。此外，還應(yīng)安裝 TorchVision，并將其作為模型合集 (model zoo)。

可通過 pip 快速安裝：

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

或參考官網(wǎng)：pytorch.org/get-started…

PyTorch 版本應(yīng)該和 TorchVision 版本兼容。

目前 TVM 支持 PyTorch 1.7 和 1.4，其他版本可能不穩(wěn)定。

import tvm
from tvm import relay
import numpy as np
from tvm.contrib.download import download_testdata
# 導(dǎo)入 PyTorch
import torch
import torchvision

加載預(yù)訓(xùn)練的 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()
輸出結(jié)果：

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]

加載測試圖像?

經(jīng)典的貓咪示例：

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))
# 預(yù)處理圖像，并將其轉(zhuǎn)換為張量
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)

將計(jì)算圖導(dǎo)入 Relay?

將 PyTorch 計(jì)算圖轉(zhuǎn)換為 Relay 計(jì)算圖。input_name 可以是任意值。

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

Relay 構(gòu)建?

用給定的輸入規(guī)范，將計(jì)算圖編譯為 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)

輸出結(jié)果：

/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 上執(zhí)行可移植計(jì)算圖?

將編譯好的模型部署到 target 上：

from tvm.contrib import graph_executor
dtype = "float32"
m = graph_executor.GraphModule(lib["default"](dev))
# 設(shè)置輸入
m.set_input(input_name, tvm.nd.array(img.astype(dtype)))
# 執(zhí)行
m.run()
# 得到輸出
tvm_output = m.get_output(0)

查找分類集名稱?

在 1000 個類的分類集中，查找分?jǐn)?shù)最高的第一個：

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 個結(jié)果
top1_tvm = np.argmax(tvm_output.numpy()[0])
tvm_class_key = class_id_to_key[top1_tvm]
# 將輸入轉(zhuǎn)換為 PyTorch 變量，并獲取 PyTorch 結(jié)果進(jìn)行比較
with torch.no_grad():
    torch_img = torch.from_numpy(img)
    output = model(torch_img)
 # 獲得 PyTorch 的前 1 個結(jié)果
    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]))

輸出結(jié)果：