Эта статья переведена с английского документа Compile PyTorch Models .
Автор Алекс Вонг .
Доступ к другим китайским документам TVM можно получить → Китайская станция TVM .
В этой статье описывается, как развернуть модели PyTorch с помощью Relay.
Сначала должен быть установлен PyTorch. Кроме того, необходимо установить TorchVision и использовать его как коллекцию моделей (зоопарк моделей).
Быстрая установка через pip:
pip install torch==1.7.0
pip install torchvision==0.8.1
Или обратитесь на официальный сайт: https://pytorch.org/get-started/locally/
Версия 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()
输出结果：
Загрузка: «https://download.pytorch.org/models/resnet18-f37072fd.pth» в /workspace/.cache/torch/hub/checkpoints/resnet18-f37072fd.pth.
0%| | 0,00/44,7M [00:00<?, ?Б/с]
11%|# | 4,87M/44,7M [00:00<00:00, 51,0 МБ/с]
22%|##1 | 9,73 М/44,7 М [00:00<00:00, 49,2 МБ/с]
74%|#######3 | 32,9M/44,7M [00:00<00:00, 136 МБ/с]
100%|##########| 44,7 млн/44,7 млн [00:00<00:00, 129 МБ/с]
Загрузите тестовое изображение
Пример классического кота:
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)
Релейная сборка
С заданной входной спецификацией скомпилируйте граф вычислений для цели llvm.
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
Разверните скомпилированную модель в цели:
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