TVM教程-----Export and Load Relax Executables

Export and Load Relax Executables目录Export and Load Relax Executables本教程演示如何将编译好的 Relax module export 为共享对象shared object再加载回 TVM runtime并在交互环境或独立脚本中运行结果。本教程展示如何使用tvm.relaxAPI将 Relax或导入的 PyTorch / ONNX程序变成可部署产物。Note本教程以 PyTorch 作为源格式但对 ONNX 模型而言export/load 工作流相同。对于 ONNX请使用from_onnx(model, keep_params_in_inputTrue)替代from_exported_program()随后按相同步骤进行 build、export 与 load。IntroductionTVM 将 Relax 程序构建为tvm.runtime.Executable对象。它们包含 VM bytecode、已编译的 kernel 以及常量。通过export_library()导出该 executable你会得到一个共享库例如 Linux 上的.so可以把它运到另一台机器、经 RPC 上传或稍后用 TVM runtime 重新加载。本教程给出端到端的确切步骤并说明过程中会产生哪些文件。importosfrompathlibimportPathtry:importtorchfromtorch.exportimportexportexceptImportError:# pragma: no covertorchNone# type: ignorePrepare a Torch MLP and Convert to Relax我们从一个小型 PyTorch MLP 开始以便示例保持轻量。模型被 export 为torch.export.ExportedProgram再翻译成 RelaxIRModule。importtvm_ffiimporttvmfromtvmimportrelaxfromtvm.relax.frontend.torchimportfrom_exported_program# Check dependencies firstIS_IN_CIos.getenv(CI,).lower()trueHAS_TORCHtorchisnotNoneRUN_EXAMPLEHAS_TORCHandnotIS_IN_CIifHAS_TORCH:classTorchMLP(torch.nn.Module):def__init__(self)-None:super().__init__()self.nettorch.nn.Sequential(torch.nn.Flatten(),torch.nn.Linear(28*28,128),torch.nn.ReLU(),torch.nn.Linear(128,10),)defforward(self,data:torch.Tensor)-torch.Tensor:# type: ignore[override]returnself.net(data)else:# pragma: no coverTorchMLPNone# type: ignore[misc, assignment]ifRUN_EXAMPLE:torch_modelTorchMLP().eval()example_args(torch.randn(1,1,28,28,dtypetorch.float32),)withtorch.no_grad():exported_programexport(torch_model,example_args)modfrom_exported_program(exported_program,keep_params_as_inputTrue)# Separate model parameters so they can be bound later (or stored on disk).mod,paramsrelax.frontend.detach_params(mod)print(Imported Relax module:)mod.show()Build and Export withexport_library我们为llvm构建以生成 CPU 代码然后导出得到的 executable。传入workspace_dir会保留中间打包文件便于检查产物内容。TARGETtvm.target.Target(llvm)ARTIFACT_DIRPath(relax_export_artifacts)ARTIFACT_DIR.mkdir(exist_okTrue)ifRUN_EXAMPLE:# Apply the default Relax compilation pipeline before building.pipelinerelax.get_pipeline()withTARGET:built_modpipeline(mod)# Build without params - well pass them at runtimeexecutabletvm.compile(built_mod,targetTARGET)library_pathARTIFACT_DIR/mlp_cpu.soexecutable.export_library(str(library_path),workspace_dirstr(ARTIFACT_DIR))print(fExported runtime library to:{library_path})# The workspace directory now contains the shared object and supporting files.produced_filessorted(p.nameforpinARTIFACT_DIR.iterdir())print(Artifacts saved:)fornameinproduced_files:print(f -{name})# Generated files:# - mlp_cpu.so: The main deployable shared library containing VM bytecode,# compiled kernels, and constants. Note: Since parameters are passed at runtime,# you will also need to save a separate parameters file (see next section).# - Intermediate object files (devc.o, lib0.o, etc.) are kept in the# workspace for inspection but are not required for deployment.## Note: Additional files like *.params, *.metadata.json, or *.imports# may appear in specific configurations but are typically embedded into the# shared library or only generated when needed.生成的文件说明mlp_cpu.so主要的可部署共享库包含 VM bytecode、已编译 kernel 和常量。注意由于参数在 runtime 传入你还需要单独保存一份参数文件见下一节。中间目标文件devc.o、lib0.o等会保留在 workspace 中供检查但部署时不需要。Note在特定配置下还可能出现*.params、*.metadata.json或*.imports等额外文件但它们通常被嵌入共享库或仅在需要时生成。Load the Exported Library and Run It一旦共享对象生成我们就可以在任意具有兼容指令集的机器上将其重新加载回 TVM runtime。Relax VM 直接消费该 runtime module。ifRUN_EXAMPLE:loaded_rt_modtvm.runtime.load_module(str(library_path))devtvm.cpu(0)vmrelax.VirtualMachine(loaded_rt_mod,dev)# Prepare input datainput_tensortorch.randn(1,1,28,28,dtypetorch.float32)vm_inputtvm.runtime.tensor(input_tensor.numpy(),dev)# Prepare parameters (allocate on target device)vm_params[tvm.runtime.tensor(p,dev)forpinparams[main]]# Run inference: pass input data followed by all parameterstvm_outputvm[main](vm_input,*vm_params)# TVM returns Array objects for tuple outputs, access via indexing.# For models imported from PyTorch, outputs are typically tuples (even for single outputs).# For ONNX models, outputs may be a single Tensor directly.ifisinstance(tvm_output,tvm_ffi.Array)andlen(tvm_output)0:result_tensortvm_output[0]else:result_tensortvm_outputprint(VM output shape:,result_tensor.shape)print(VM output type:,type(tvm_output),-,type(result_tensor))# You can still inspect the executable after reloading.print(Executable stats:\n,loaded_rt_mod[stats]())Save Parameters for Deployment由于参数在 runtime 传入并未嵌入.so我们必须单独保存它们以便部署。这是在其他机器或独立脚本中使用该模型的必要步骤。importnumpyasnpifRUN_EXAMPLE:# Save parameters to diskparams_pathARTIFACT_DIR/model_params.npzparam_arrays{fp_{i}:p.numpy()fori,pinenumerate(params[main])}np.savez(str(params_path),**param_arrays)print(fSaved parameters to:{params_path})# Note: Alternatively, you can embed parameters directly into the .so to# create a single-file deployment. Use keep_params_as_inputFalse when# importing from PyTorch:## .. code-block:: python## mod from_exported_program(exported_program, keep_params_as_inputFalse)# # Parameters are now embedded as constants in the module# executable tvm.compile(built_mod, targetTARGET)# # Runtime: vm[main](input) # No need to pass params!## This creates a single-file deployment (only the .so is needed), but you# lose the flexibility to swap parameters without recompiling. For most# production workflows, separating code and parameters (as shown above) is# preferred for flexibility.Note或者你也可以把参数直接嵌入.so形成单文件部署。从 PyTorch 导入时使用keep_params_as_inputFalsemodfrom_exported_program(exported_program,keep_params_as_inputFalse)# Parameters are now embedded as constants in the moduleexecutabletvm.compile(built_mod,targetTARGET)# Runtime: vm[main](input) # No need to pass params!这样只需.so一个文件即可部署但会失去不重新编译就更换参数的灵活性。对多数生产工作流而言按上文所示将代码与参数分离通常更灵活、更可取。Loading and Running the Exported Model要在另一台机器或独立脚本中使用已导出的模型你需要同时加载.so库和参数文件。下面是重新加载并运行模型的完整示例。将其保存为run_mlp.py要从命令行直接执行chmodx run_mlp.py ./run_mlp.py# Run it like a regular program完整脚本#!/usr/bin/env python3importnumpyasnpimporttvmfromtvmimportrelax# Step 1: Load the compiled librarylibtvm.runtime.load_module(relax_export_artifacts/mlp_cpu.so)# Step 2: Create Virtual Machinedevicetvm.cpu(0)vmrelax.VirtualMachine(lib,device)# Step 3: Load parameters from the .npz fileparams_npznp.load(relax_export_artifacts/model_params.npz)params[tvm.runtime.tensor(params_npz[fp_{i}],device)foriinrange(len(params_npz))]# Step 4: Prepare input datadatanp.random.randn(1,1,28,28).astype(float32)input_tensortvm.runtime.tensor(data,device)# Step 5: Run inference (pass input followed by all parameters)outputvm[main](input_tensor,*params)# Step 6: Extract result (output may be tuple or single Tensor)# PyTorch models typically return tuples, ONNX models may return a single Tensorifisinstance(output,tvm_ffi.Array)andlen(output)0:result_tensoroutput[0]else:result_tensoroutputprint(Prediction shape:,result_tensor.shape)print(Predicted class:,np.argmax(result_tensor.numpy()))在 GPU 上运行若要在 GPU 而非 CPU 上运行做如下修改为 GPU 编译教程前文大约在 line 112 附近TARGETtvm.target.Target(cuda)# Change from llvm to cuda在脚本中使用 GPU devicedevicetvm.cuda(0)# Use CUDA device instead of CPUvmrelax.VirtualMachine(lib,device)# Load parameters to GPUparams[tvm.runtime.tensor(params_npz[fp_{i}],device)# Note: device parameterforiinrange(len(params_npz))]# Prepare input on GPUinput_tensortvm.runtime.tensor(data,device)# Note: device parameter脚本其余部分保持不变。所有张量参数与输入都必须分配在与已编译模型相同的设备GPU上。部署检查清单当迁移到另一台主机经 RPC 或 SCP时你必须复制两个文件mlp_cpu.soGPU 则为mlp_cuda.so——已编译的模型代码model_params.npz——以 NumPy 数组序列化的模型参数远端机器需要这两个文件位于同一目录。上面的脚本假设它们相对于脚本位置在relax_export_artifacts/下。请按你的部署调整路径。对于 GPU 部署请确保目标机器有兼容的 CUDA 驱动且模型是为相同的 GPU 架构编译的。Deploying to Remote Devices要将已导出的模型部署到远端 ARM Linux 设备例如 Raspberry Pi可以使用 TVM 的 RPC 机制进行交叉编译、上传并远程运行模型。该工作流在以下场景很有用目标设备编译资源有限你希望在真实硬件上运行以微调性能你需要部署到嵌入式设备参见 cross_compilation_and_rpc 获取完整指南涵盖在远端设备上设置 TVM runtime在设备上启动 RPC server为 ARM 目标交叉编译例如llvm -mtripleaarch64-linux-gnu经 RPC 上传已导出的库远程运行推理ARM 部署工作流的快速示例importtvm.rpcasrpcfromtvmimportrelax# Step 1: Cross-compile for ARM target (on local machine)TARGETtvm.target.Target({kind:llvm,mtriple:aarch64-linux-gnu})executabletvm.compile(built_mod,targetTARGET)executable.export_library(mlp_arm.so)# Step 2: Connect to remote device RPC serverremoterpc.connect(192.168.1.100,9090)# Device IP and RPC port# Step 3: Upload the compiled library and parametersremote.upload(mlp_arm.so)remote.upload(model_params.npz)# Step 4: Load and run on remote devicelibremote.load_module(mlp_arm.so)vmrelax.VirtualMachine(lib,remote.cpu())# ... prepare input and params, then run inference关键差异在于编译时使用 ARM target triple并通过 RPC 上传文件而不是直接拷贝。FAQCan I run the.soas a standalone executable (like./mlp_cpu.so)?不行。.so文件是共享库不是独立可执行二进制。你不能从终端直接运行它。必须通过 TVM runtime 程序加载如上文 “Loading and Running” 一节所示。.so打包了 VM bytecode 与已编译 kernel但仍需要 TVM runtime 才能执行。Which devices can run the exported library?目标必须与你编译时所针对的 ISA 匹配本例为llvm。只要 target triple、runtime ABI 以及可用设备对齐就可以在机器之间迁移该产物。对于异构构建CPU 加 GPU还需要一并携带额外的设备库。What about the.paramsandmetadata.jsonfiles?这些辅助文件仅在特定配置下生成。在本教程中由于我们在 runtime 传入参数它们不会生成。当它们确实出现时可以与.so放在一起供检查但核心内容通常已嵌入共享对象本身因此通常仅部署.so就足够了。Download Jupyter notebook: export_and_load_executable.ipynbDownload Python source code: export_and_load_executable.py