|
| 1 | +# Licensed to the Apache Software Foundation (ASF) under one |
| 2 | +# or more contributor license agreements. See the NOTICE file |
| 3 | +# distributed with this work for additional information |
| 4 | +# regarding copyright ownership. The ASF licenses this file |
| 5 | +# to you under the Apache License, Version 2.0 (the |
| 6 | +# "License"); you may not use this file except in compliance |
| 7 | +# with the License. You may obtain a copy of the License at |
| 8 | +# |
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +# |
| 11 | +# Unless required by applicable law or agreed to in writing, |
| 12 | +# software distributed under the License is distributed on an |
| 13 | +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 14 | +# KIND, either express or implied. See the License for the |
| 15 | +# specific language governing permissions and limitations |
| 16 | +# under the License. |
| 17 | + |
| 18 | +""" |
| 19 | +.. _optimize_model: |
| 20 | +
|
| 21 | +End-to-End Optimize Model |
| 22 | +========================= |
| 23 | +This tutorial demonstrates how to optimize a machine learning model using Apache TVM. We will |
| 24 | +use a pre-trained ResNet-18 model from PyTorch and end-to-end optimize it using TVM's Relax API. |
| 25 | +Please note that default end-to-end optimization may not suit complex models. |
| 26 | +""" |
| 27 | + |
| 28 | +###################################################################### |
| 29 | +# Preparation |
| 30 | +# ----------- |
| 31 | +# First, we prepare the model and input information. We use a pre-trained ResNet-18 model from |
| 32 | +# PyTorch. |
| 33 | + |
| 34 | +import os |
| 35 | +import sys |
| 36 | +import numpy as np |
| 37 | +import torch |
| 38 | +from torch import fx |
| 39 | +from torchvision.models.resnet import ResNet18_Weights, resnet18 |
| 40 | + |
| 41 | +torch_model = resnet18(weights=ResNet18_Weights.DEFAULT) |
| 42 | + |
| 43 | +###################################################################### |
| 44 | +# Review Overall Flow |
| 45 | +# ------------------- |
| 46 | +# .. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/design/tvm_overall_flow.svg |
| 47 | +# :align: center |
| 48 | +# :width: 80% |
| 49 | +# |
| 50 | +# The overall flow consists of the following steps: |
| 51 | +# |
| 52 | +# - **Construct or Import a Model**: Construct a neural network model or import a pre-trained |
| 53 | +# model from other frameworks (e.g. PyTorch, ONNX), and create the TVM IRModule, which contains |
| 54 | +# all the information needed for compilation, including high-level Relax functions for |
| 55 | +# computational graph, and low-level TensorIR functions for tensor program. |
| 56 | +# - **Perform Composable Optimizations**: Perform a series of optimization transformations, |
| 57 | +# such as graph optimizations, tensor program optimizations, and library dispatching. |
| 58 | +# - **Build and Universal Deployment**: Build the optimized model to a deployable module to the |
| 59 | +# universal runtime, and execute it on different devices, such as CPU, GPU, or other accelerators. |
| 60 | +# |
| 61 | + |
| 62 | + |
| 63 | +###################################################################### |
| 64 | +# Convert the model to IRModule |
| 65 | +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ |
| 66 | +# Next step, we convert the model to an IRModule using the Relax frontend for PyTorch for further |
| 67 | +# optimization. Besides the model, we also need to provide the input shape and data type. |
| 68 | + |
| 69 | +import tvm |
| 70 | +from tvm import relax |
| 71 | +from tvm.relax.frontend.torch import from_fx |
| 72 | + |
| 73 | +torch_model = resnet18(weights=ResNet18_Weights.DEFAULT) |
| 74 | + |
| 75 | +# Give the input shape and data type |
| 76 | +input_info = [((1, 3, 224, 224), "float32")] |
| 77 | + |
| 78 | +# Convert the model to IRModule |
| 79 | +with torch.no_grad(): |
| 80 | + torch_fx_model = fx.symbolic_trace(torch_model) |
| 81 | + mod = from_fx(torch_fx_model, input_info, keep_params_as_input=True) |
| 82 | + |
| 83 | +mod, params = relax.frontend.detach_params(mod) |
| 84 | +mod.show() |
| 85 | + |
| 86 | +###################################################################### |
| 87 | +# IRModule Optimization |
| 88 | +# --------------------- |
| 89 | +# Apache TVM Unity provides a flexible way to optimize the IRModule. Everything centered |
| 90 | +# around IRModule optimization can be composed with existing pipelines. Note that each |
| 91 | +# transformation can be combined as an optimization pipeline via ``tvm.ir.transform.Sequential``. |
| 92 | +# |
| 93 | +# In this tutorial, we focus on the end-to-end optimization of the model via auto-tuning. We |
| 94 | +# leverage MetaSchedule to tune the model and store the tuning logs to the database. We also |
| 95 | +# apply the database to the model to get the best performance. |
| 96 | +# |
| 97 | + |
| 98 | +TOTAL_TRIALS = 8000 # Change to 20000 for better performance if needed |
| 99 | +target = tvm.target.Target("nvidia/geforce-rtx-3090-ti") # Change to your target device |
| 100 | +work_dir = "tuning_logs" |
| 101 | + |
| 102 | +# Skip running in CI environment |
| 103 | +IS_IN_CI = os.getenv("CI", "") == "true" |
| 104 | +if IS_IN_CI: |
| 105 | + sys.exit(0) |
| 106 | + |
| 107 | +with target: |
| 108 | + mod = tvm.ir.transform.Sequential( |
| 109 | + [ |
| 110 | + # Convert BatchNorm into a sequence of simpler ops for fusion |
| 111 | + relax.transform.DecomposeOpsForInference(), |
| 112 | + # Canonicalize the bindings |
| 113 | + relax.transform.CanonicalizeBindings(), |
| 114 | + # Run default optimization pipeline |
| 115 | + relax.get_pipeline("zero"), |
| 116 | + # Tune the model and store the log to database |
| 117 | + relax.transform.MetaScheduleTuneIRMod({}, work_dir, TOTAL_TRIALS), |
| 118 | + # Apply the database |
| 119 | + relax.transform.MetaScheduleApplyDatabase(work_dir), |
| 120 | + ] |
| 121 | + )(mod) |
| 122 | + |
| 123 | +# Only show the main function |
| 124 | +mod["main"].show() |
| 125 | + |
| 126 | +###################################################################### |
| 127 | +# Build and Deploy |
| 128 | +# ---------------- |
| 129 | +# Finally, we build the optimized model and deploy it to the target device. |
| 130 | + |
| 131 | +ex = relax.build(mod, target="cuda") |
| 132 | +dev = tvm.device("cuda", 0) |
| 133 | +vm = relax.VirtualMachine(ex, dev) |
| 134 | +# Need to allocate data and params on GPU device |
| 135 | +gpu_data = tvm.nd.array(np.random.rand(1, 3, 224, 224).astype("float32"), dev) |
| 136 | +gpu_params = [tvm.nd.array(p, dev) for p in params["main"]] |
| 137 | +gpu_out = vm["main"](gpu_data, *gpu_params).numpy() |
| 138 | + |
| 139 | +print(gpu_out.shape) |
0 commit comments