{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(sphx_glr_how_to_compile_models_from_pytorch.py)=\n",
"# 编译 PyTorch 模型\n",
"\n",
"**Author**: [Alex Wong](https://github.com/alexwong/)\n",
"\n",
"本文是使用 Relay 部署 PyTorch 模型的入门教程。\n",
"\n",
"对于我们来说，首先应该安装 PyTorch。TorchVision 也是必需的，因为我们将使用它作为我们的模型动物园。\n",
"\n",
"快速的解决方案是通过 pip 进行安装：\n",
"\n",
"```python\n",
"pip install torch==1.7.0\n",
"pip install torchvision==0.8.1\n",
"```\n",
"\n",
"或者参考[官方网站](https://pytorch.org/get-started/locally/)。\n",
"\n",
"PyTorch 版本应该向后兼容，但应该与适当的 TorchVision 版本一起使用。\n",
"\n",
"目前，TVM 支持 PyTorch 1.7 和 1.4。其他版本可能不稳定。"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"# PyTorch imports\n",
"import torch\n",
"import torchvision\n",
"\n",
"import set_env # 设置 TVM 环境\n",
"import tvm\n",
"from tvm import relay\n",
"from tvm.contrib.download import download_testdata"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 载入 PyTorch 预训练模型"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"model_name = \"resnet18\"\n",
"model = getattr(torchvision.models, model_name)(pretrained=True)\n",
"model = model.eval()\n",
"\n",
"# 我们通过跟踪获取 TorchScripted 模型\n",
"input_shape = [1, 3, 224, 224]\n",
"input_data = torch.randn(input_shape)\n",
"scripted_model = torch.jit.trace(model, input_data).eval()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 加载测试图片"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"from PIL import Image\n",
"\n",
"img_url = \"https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true\"\n",
"img_path = download_testdata(img_url, \"cat.png\", module=\"data\")\n",
"img = Image.open(img_path).resize((224, 224))\n",
"\n",
"# Preprocess the image and convert to tensor\n",
"from torchvision import transforms\n",
"\n",
"my_preprocess = transforms.Compose(\n",
" [\n",
" transforms.Resize(256),\n",
" transforms.CenterCrop(224),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
" ]\n",
")\n",
"img = my_preprocess(img)\n",
"img = np.expand_dims(img, 0)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 导入 Graph 到 Relay\n",
"\n",
"将 PyTorch 图转换为 Relay 图。`input_name` 可以是任意的。"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"input_name = \"input0\"\n",
"shape_list = [(input_name, img.shape)]\n",
"mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Relay 构建\n",
"\n",
"使用给定的输入规范将 graph 编译为 llvm 目标："
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.\n"
]
}
],
"source": [
"target = tvm.target.Target(\"llvm\", host=\"llvm\")\n",
"dev = tvm.cpu(0)\n",
"with tvm.transform.PassContext(opt_level=3):\n",
" lib = relay.build(mod, target=target, params=params)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 在 TVM 上执行可移植 Graph\n",
"\n",
"可以尝试在目标上部署编译后的模型。"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"from tvm.contrib import graph_executor\n",
"\n",
"dtype = \"float32\"\n",
"m = graph_executor.GraphModule(lib[\"default\"](dev))\n",
"# Set inputs\n",
"m.set_input(input_name, tvm.nd.array(img.astype(dtype)))\n",
"# Execute\n",
"m.run()\n",
"# Get outputs\n",
"tvm_output = m.get_output(0)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 查找 synset 名称\n",
"\n",
"在 1000 类 synset 中查找预测 top 1 索引。"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Relay top-1 id: 281, class name: tabby, tabby cat\n",
"Torch top-1 id: 281, class name: tabby, tabby cat\n"
]
}
],
"source": [
"synset_url = \"\".join(\n",
" [\n",
" \"https://raw.githubusercontent.com/Cadene/\",\n",
" \"pretrained-models.pytorch/master/data/\",\n",
" \"imagenet_synsets.txt\",\n",
" ]\n",
")\n",
"synset_name = \"imagenet_synsets.txt\"\n",
"synset_path = download_testdata(synset_url, synset_name, module=\"data\")\n",
"with open(synset_path) as f:\n",
" synsets = f.readlines()\n",
"\n",
"synsets = [x.strip() for x in synsets]\n",
"splits = [line.split(\" \") for line in synsets]\n",
"key_to_classname = {spl[0]: \" \".join(spl[1:]) for spl in splits}\n",
"\n",
"class_url = \"\".join(\n",
" [\n",
" \"https://raw.githubusercontent.com/Cadene/\",\n",
" \"pretrained-models.pytorch/master/data/\",\n",
" \"imagenet_classes.txt\",\n",
" ]\n",
")\n",
"class_name = \"imagenet_classes.txt\"\n",
"class_path = download_testdata(class_url, class_name, module=\"data\")\n",
"with open(class_path) as f:\n",
" class_id_to_key = f.readlines()\n",
"\n",
"class_id_to_key = [x.strip() for x in class_id_to_key]\n",
"\n",
"# Get top-1 result for TVM\n",
"top1_tvm = np.argmax(tvm_output.numpy()[0])\n",
"tvm_class_key = class_id_to_key[top1_tvm]\n",
"\n",
"# Convert input to PyTorch variable and get PyTorch result for comparison\n",
"with torch.no_grad():\n",
" torch_img = torch.from_numpy(img)\n",
" output = model(torch_img)\n",
"\n",
" # Get top-1 result for PyTorch\n",
" top1_torch = np.argmax(output.numpy())\n",
" torch_class_key = class_id_to_key[top1_torch]\n",
"\n",
"print(\"Relay top-1 id: {}, class name: {}\".format(top1_tvm, key_to_classname[tvm_class_key]))\n",
"print(\"Torch top-1 id: {}, class name: {}\".format(top1_torch, key_to_classname[torch_class_key]))"
]
}
],
"metadata": {
"interpreter": {
"hash": "6ee5142ba8a2589df39b0df03e82f50c3ae535c49aaf7d83abad1a0d572c7e37"
},
"kernelspec": {
"display_name": "Python 3.10.4 ('tvm-test')",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.4"
},
"orig_nbformat": 4
},
"nbformat": 4,
"nbformat_minor": 2
}