You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "tlopex (via GitHub)" <gi...@apache.org> on 2025/11/01 19:54:13 UTC
[PR] [Docs][Relax] Update cross-compilation and RPC tutorial with modern PyTorch deployment workflow [tvm]
tlopex opened a new pull request, #18413:
URL: https://github.com/apache/tvm/pull/18413
This PR modernizes the cross-compilation and RPC tutorial by adding a complete PyTorch/Relax deployment workflow alongside the existing TE examples.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [DOCS] Update cross-compilation and RPC tutorial with modern PyTorch deployment workflow [tvm]
Posted by "tlopex (via GitHub)" <gi...@apache.org>.
tlopex commented on PR #18413:
URL: https://github.com/apache/tvm/pull/18413#issuecomment-3478963754
cc @tqchen
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [DOCS] Update cross-compilation and RPC tutorial with modern PyTorch deployment workflow [tvm]
Posted by "MasterJH5574 (via GitHub)" <gi...@apache.org>.
MasterJH5574 merged PR #18413:
URL: https://github.com/apache/tvm/pull/18413
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [DOCS] Update cross-compilation and RPC tutorial with modern PyTorch deployment workflow [tvm]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #18413:
URL: https://github.com/apache/tvm/pull/18413#discussion_r2483892546
##########
docs/how_to/tutorials/cross_compilation_and_rpc.py:
##########
@@ -256,13 +256,328 @@ def run_opencl():
print("OpenCL test passed!")
+#########################################################################
+# Deploy PyTorch Models to Remote Devices with RPC
+# ------------------------------------------------
+# The above examples demonstrate cross compilation and RPC using low-level
+# TensorIR (via TE). For deploying complete neural network models from frameworks
+# like PyTorch or ONNX, TVM's Relax provides a higher-level abstraction that is
+# better suited for end-to-end model compilation.
+#
+# This section shows a modern workflow for deploying models to **any remote device**:
+#
+# 1. Import a PyTorch model and convert it to Relax
+# 2. Cross-compile for the target architecture (ARM, x86, RISC-V, etc.)
+# 3. Deploy via RPC to a remote device
+# 4. Run inference remotely
+#
+# This workflow is applicable to various deployment scenarios:
+#
+# - **ARM devices**: Raspberry Pi, NVIDIA Jetson, mobile phones
+# - **x86 servers**: Remote Linux servers, cloud instances
+# - **Embedded systems**: RISC-V boards, custom hardware
+# - **Accelerators**: Remote machines with GPUs, TPUs, or other accelerators
+#
+# .. note::
+# This example uses PyTorch for demonstration, but the workflow is identical
+# for ONNX models. Simply replace ``from_exported_program()`` with
+# ``from_onnx(model, keep_params_in_input=True)`` and follow the same steps.
+
+# First, let's check if PyTorch is available
+try:
+ import torch
+ from torch.export import export
+
+ HAS_TORCH = True
+except ImportError:
+ HAS_TORCH = False
+ print("PyTorch not available, skipping Relax example")
+
+
+def run_pytorch_model_via_rpc():
+ """
+ Demonstrates the complete workflow of deploying a PyTorch model to an ARM device via RPC.
+ """
+ if not HAS_TORCH:
+ print("Skipping PyTorch example (PyTorch not installed)")
+ return
+
+ from tvm import relax
+ from tvm.relax.frontend.torch import from_exported_program
+
+ ######################################################################
+ # Step 1: Define and Export PyTorch Model
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # We use a simple MLP model for demonstration. In practice, this could be
+ # any PyTorch model (ResNet, BERT, etc.).
+
+ class TorchMLP(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.net = torch.nn.Sequential(
+ torch.nn.Flatten(),
+ torch.nn.Linear(28 * 28, 128),
+ torch.nn.ReLU(),
+ torch.nn.Linear(128, 10),
+ )
+
+ def forward(self, data: torch.Tensor) -> torch.Tensor:
+ return self.net(data)
+
+ # Export the model using PyTorch 2.x export API
+ torch_model = TorchMLP().eval()
+ example_args = (torch.randn(1, 1, 28, 28, dtype=torch.float32),)
+
+ with torch.no_grad():
+ exported_program = export(torch_model, example_args)
+
+ ######################################################################
+ # Step 2: Convert to Relax and Prepare for Compilation
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Convert the exported PyTorch program to TVM's Relax representation
+
+ mod = from_exported_program(exported_program, keep_params_as_input=True)
+ # Separate parameters from the model for flexible deployment
+ mod, params = relax.frontend.detach_params(mod)
+
+ print("Converted PyTorch model to Relax:")
+ print(f" - Number of parameters: {len(params['main'])}")
+
+ ######################################################################
+ # Step 3: Cross-Compile for Target Device
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Compile the model for the target device architecture. The target
+ # configuration depends on your deployment scenario.
+
+ if local_demo:
+ # For demonstration on local machine, use local target
+ target = tvm.target.Target("llvm")
+ print("Using local target for demonstration")
+ else:
+ # Choose the appropriate target for your device:
+ #
+ # ARM devices:
+ # - Raspberry Pi 3/4 (32-bit): "llvm -mtriple=armv7l-linux-gnueabihf"
+ # - Raspberry Pi 4 (64-bit) / Jetson: "llvm -mtriple=aarch64-linux-gnu"
+ # - Android: "llvm -mtriple=aarch64-linux-android"
+ #
+ # x86 servers:
+ # - Linux x86_64: "llvm -mtriple=x86_64-linux-gnu"
+ # - With AVX-512: "llvm -mtriple=x86_64-linux-gnu -mcpu=skylake-avx512"
+ #
+ # RISC-V:
+ # - RV64: "llvm -mtriple=riscv64-unknown-linux-gnu"
+ #
+ # GPU targets:
+ # - CUDA: tvm.target.Target("cuda", host="llvm -mtriple=x86_64-linux-gnu")
+ # - OpenCL: tvm.target.Target("opencl", host="llvm -mtriple=aarch64-linux-gnu")
+ #
+ # For this example, we use ARM 64-bit
+ target = tvm.target.Target("llvm -mtriple=aarch64-linux-gnu")
+ print(f"Cross-compiling for target: {target}")
+
+ # Apply optimization pipeline
+ pipeline = relax.get_pipeline()
+ with target:
+ built_mod = pipeline(mod)
+
+ # Compile to executable
+ executable = tvm.compile(built_mod, target=target)
+
+ # Export to shared library
+ lib_path = temp.relpath("model_deployed.so")
+ executable.export_library(lib_path)
+ print(f"Exported library to: {lib_path}")
+
+ # Save parameters separately
+ import numpy as np
Review Comment:
The `numpy` library is already imported as `np` at the top of the file (line 97). This import is redundant and can be removed to keep the code clean.
##########
docs/how_to/tutorials/cross_compilation_and_rpc.py:
##########
@@ -256,13 +256,328 @@ def run_opencl():
print("OpenCL test passed!")
+#########################################################################
+# Deploy PyTorch Models to Remote Devices with RPC
+# ------------------------------------------------
+# The above examples demonstrate cross compilation and RPC using low-level
+# TensorIR (via TE). For deploying complete neural network models from frameworks
+# like PyTorch or ONNX, TVM's Relax provides a higher-level abstraction that is
+# better suited for end-to-end model compilation.
+#
+# This section shows a modern workflow for deploying models to **any remote device**:
+#
+# 1. Import a PyTorch model and convert it to Relax
+# 2. Cross-compile for the target architecture (ARM, x86, RISC-V, etc.)
+# 3. Deploy via RPC to a remote device
+# 4. Run inference remotely
+#
+# This workflow is applicable to various deployment scenarios:
+#
+# - **ARM devices**: Raspberry Pi, NVIDIA Jetson, mobile phones
+# - **x86 servers**: Remote Linux servers, cloud instances
+# - **Embedded systems**: RISC-V boards, custom hardware
+# - **Accelerators**: Remote machines with GPUs, TPUs, or other accelerators
+#
+# .. note::
+# This example uses PyTorch for demonstration, but the workflow is identical
+# for ONNX models. Simply replace ``from_exported_program()`` with
+# ``from_onnx(model, keep_params_in_input=True)`` and follow the same steps.
+
+# First, let's check if PyTorch is available
+try:
+ import torch
+ from torch.export import export
+
+ HAS_TORCH = True
+except ImportError:
+ HAS_TORCH = False
+ print("PyTorch not available, skipping Relax example")
+
+
+def run_pytorch_model_via_rpc():
+ """
+ Demonstrates the complete workflow of deploying a PyTorch model to an ARM device via RPC.
+ """
+ if not HAS_TORCH:
+ print("Skipping PyTorch example (PyTorch not installed)")
+ return
+
+ from tvm import relax
+ from tvm.relax.frontend.torch import from_exported_program
+
+ ######################################################################
+ # Step 1: Define and Export PyTorch Model
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # We use a simple MLP model for demonstration. In practice, this could be
+ # any PyTorch model (ResNet, BERT, etc.).
+
+ class TorchMLP(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.net = torch.nn.Sequential(
+ torch.nn.Flatten(),
+ torch.nn.Linear(28 * 28, 128),
+ torch.nn.ReLU(),
+ torch.nn.Linear(128, 10),
+ )
+
+ def forward(self, data: torch.Tensor) -> torch.Tensor:
+ return self.net(data)
+
+ # Export the model using PyTorch 2.x export API
+ torch_model = TorchMLP().eval()
+ example_args = (torch.randn(1, 1, 28, 28, dtype=torch.float32),)
+
+ with torch.no_grad():
+ exported_program = export(torch_model, example_args)
+
+ ######################################################################
+ # Step 2: Convert to Relax and Prepare for Compilation
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Convert the exported PyTorch program to TVM's Relax representation
+
+ mod = from_exported_program(exported_program, keep_params_as_input=True)
+ # Separate parameters from the model for flexible deployment
+ mod, params = relax.frontend.detach_params(mod)
+
+ print("Converted PyTorch model to Relax:")
+ print(f" - Number of parameters: {len(params['main'])}")
+
+ ######################################################################
+ # Step 3: Cross-Compile for Target Device
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Compile the model for the target device architecture. The target
+ # configuration depends on your deployment scenario.
+
+ if local_demo:
+ # For demonstration on local machine, use local target
+ target = tvm.target.Target("llvm")
+ print("Using local target for demonstration")
+ else:
+ # Choose the appropriate target for your device:
+ #
+ # ARM devices:
+ # - Raspberry Pi 3/4 (32-bit): "llvm -mtriple=armv7l-linux-gnueabihf"
+ # - Raspberry Pi 4 (64-bit) / Jetson: "llvm -mtriple=aarch64-linux-gnu"
+ # - Android: "llvm -mtriple=aarch64-linux-android"
+ #
+ # x86 servers:
+ # - Linux x86_64: "llvm -mtriple=x86_64-linux-gnu"
+ # - With AVX-512: "llvm -mtriple=x86_64-linux-gnu -mcpu=skylake-avx512"
+ #
+ # RISC-V:
+ # - RV64: "llvm -mtriple=riscv64-unknown-linux-gnu"
+ #
+ # GPU targets:
+ # - CUDA: tvm.target.Target("cuda", host="llvm -mtriple=x86_64-linux-gnu")
+ # - OpenCL: tvm.target.Target("opencl", host="llvm -mtriple=aarch64-linux-gnu")
+ #
+ # For this example, we use ARM 64-bit
+ target = tvm.target.Target("llvm -mtriple=aarch64-linux-gnu")
+ print(f"Cross-compiling for target: {target}")
+
+ # Apply optimization pipeline
+ pipeline = relax.get_pipeline()
+ with target:
+ built_mod = pipeline(mod)
+
+ # Compile to executable
+ executable = tvm.compile(built_mod, target=target)
+
+ # Export to shared library
+ lib_path = temp.relpath("model_deployed.so")
+ executable.export_library(lib_path)
+ print(f"Exported library to: {lib_path}")
+
+ # Save parameters separately
+ import numpy as np
+
+ params_path = temp.relpath("model_params.npz")
+ param_arrays = {f"p_{i}": p.numpy() for i, p in enumerate(params["main"])}
+ np.savez(params_path, **param_arrays)
+ print(f"Saved parameters to: {params_path}")
+
+ ######################################################################
+ # Step 4: Deploy to Remote Device via RPC
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Connect to the remote device, upload the compiled library and parameters,
+ # then run inference remotely. This works for any device with TVM RPC server.
+ #
+ # Note: The following code demonstrates the RPC workflow. In local_demo mode,
+ # we skip actual execution to avoid LocalSession compatibility issues.
+
+ if local_demo:
+ # For demonstration, show the code structure without execution
+ print("\nRPC workflow (works for any remote device):")
+ print("=" * 50)
+ print("1. Start RPC server on target device:")
+ print(" python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090")
+ print("\n2. Connect from local machine:")
+ print(" remote = rpc.connect('DEVICE_IP', 9090)")
+ print("\n3. Upload compiled library:")
+ print(" remote.upload('model_deployed.so')")
+ print(" remote.upload('model_params.npz')")
+ print("\n4. Load and run remotely:")
+ print(" lib = remote.load_module('model_deployed.so')")
+ print(" vm = relax.VirtualMachine(lib, remote.cpu())")
+ print(" result = vm['main'](input, *params)")
Review Comment:
The example code string uses `input` as a variable name. Since `input` is a Python built-in function, using it as a variable name can be confusing for readers. It would be clearer to use a different name, such as `input_data`, which is consistent with the variable used later in the actual execution part of this tutorial (line 469).
```suggestion
print(" result = vm['main'](input_data, *params)")
```
##########
docs/how_to/tutorials/cross_compilation_and_rpc.py:
##########
@@ -256,13 +256,328 @@ def run_opencl():
print("OpenCL test passed!")
+#########################################################################
+# Deploy PyTorch Models to Remote Devices with RPC
+# ------------------------------------------------
+# The above examples demonstrate cross compilation and RPC using low-level
+# TensorIR (via TE). For deploying complete neural network models from frameworks
+# like PyTorch or ONNX, TVM's Relax provides a higher-level abstraction that is
+# better suited for end-to-end model compilation.
+#
+# This section shows a modern workflow for deploying models to **any remote device**:
+#
+# 1. Import a PyTorch model and convert it to Relax
+# 2. Cross-compile for the target architecture (ARM, x86, RISC-V, etc.)
+# 3. Deploy via RPC to a remote device
+# 4. Run inference remotely
+#
+# This workflow is applicable to various deployment scenarios:
+#
+# - **ARM devices**: Raspberry Pi, NVIDIA Jetson, mobile phones
+# - **x86 servers**: Remote Linux servers, cloud instances
+# - **Embedded systems**: RISC-V boards, custom hardware
+# - **Accelerators**: Remote machines with GPUs, TPUs, or other accelerators
+#
+# .. note::
+# This example uses PyTorch for demonstration, but the workflow is identical
+# for ONNX models. Simply replace ``from_exported_program()`` with
+# ``from_onnx(model, keep_params_in_input=True)`` and follow the same steps.
+
+# First, let's check if PyTorch is available
+try:
+ import torch
+ from torch.export import export
+
+ HAS_TORCH = True
+except ImportError:
+ HAS_TORCH = False
+ print("PyTorch not available, skipping Relax example")
+
+
+def run_pytorch_model_via_rpc():
+ """
+ Demonstrates the complete workflow of deploying a PyTorch model to an ARM device via RPC.
+ """
+ if not HAS_TORCH:
+ print("Skipping PyTorch example (PyTorch not installed)")
+ return
+
+ from tvm import relax
+ from tvm.relax.frontend.torch import from_exported_program
+
+ ######################################################################
+ # Step 1: Define and Export PyTorch Model
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # We use a simple MLP model for demonstration. In practice, this could be
+ # any PyTorch model (ResNet, BERT, etc.).
+
+ class TorchMLP(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.net = torch.nn.Sequential(
+ torch.nn.Flatten(),
+ torch.nn.Linear(28 * 28, 128),
+ torch.nn.ReLU(),
+ torch.nn.Linear(128, 10),
+ )
+
+ def forward(self, data: torch.Tensor) -> torch.Tensor:
+ return self.net(data)
+
+ # Export the model using PyTorch 2.x export API
+ torch_model = TorchMLP().eval()
+ example_args = (torch.randn(1, 1, 28, 28, dtype=torch.float32),)
+
+ with torch.no_grad():
+ exported_program = export(torch_model, example_args)
+
+ ######################################################################
+ # Step 2: Convert to Relax and Prepare for Compilation
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Convert the exported PyTorch program to TVM's Relax representation
+
+ mod = from_exported_program(exported_program, keep_params_as_input=True)
+ # Separate parameters from the model for flexible deployment
+ mod, params = relax.frontend.detach_params(mod)
+
+ print("Converted PyTorch model to Relax:")
+ print(f" - Number of parameters: {len(params['main'])}")
+
+ ######################################################################
+ # Step 3: Cross-Compile for Target Device
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Compile the model for the target device architecture. The target
+ # configuration depends on your deployment scenario.
+
+ if local_demo:
+ # For demonstration on local machine, use local target
+ target = tvm.target.Target("llvm")
+ print("Using local target for demonstration")
+ else:
+ # Choose the appropriate target for your device:
+ #
+ # ARM devices:
+ # - Raspberry Pi 3/4 (32-bit): "llvm -mtriple=armv7l-linux-gnueabihf"
+ # - Raspberry Pi 4 (64-bit) / Jetson: "llvm -mtriple=aarch64-linux-gnu"
+ # - Android: "llvm -mtriple=aarch64-linux-android"
+ #
+ # x86 servers:
+ # - Linux x86_64: "llvm -mtriple=x86_64-linux-gnu"
+ # - With AVX-512: "llvm -mtriple=x86_64-linux-gnu -mcpu=skylake-avx512"
+ #
+ # RISC-V:
+ # - RV64: "llvm -mtriple=riscv64-unknown-linux-gnu"
+ #
+ # GPU targets:
+ # - CUDA: tvm.target.Target("cuda", host="llvm -mtriple=x86_64-linux-gnu")
+ # - OpenCL: tvm.target.Target("opencl", host="llvm -mtriple=aarch64-linux-gnu")
+ #
+ # For this example, we use ARM 64-bit
+ target = tvm.target.Target("llvm -mtriple=aarch64-linux-gnu")
+ print(f"Cross-compiling for target: {target}")
+
+ # Apply optimization pipeline
+ pipeline = relax.get_pipeline()
+ with target:
+ built_mod = pipeline(mod)
+
+ # Compile to executable
+ executable = tvm.compile(built_mod, target=target)
+
+ # Export to shared library
+ lib_path = temp.relpath("model_deployed.so")
+ executable.export_library(lib_path)
+ print(f"Exported library to: {lib_path}")
+
+ # Save parameters separately
+ import numpy as np
+
+ params_path = temp.relpath("model_params.npz")
+ param_arrays = {f"p_{i}": p.numpy() for i, p in enumerate(params["main"])}
+ np.savez(params_path, **param_arrays)
+ print(f"Saved parameters to: {params_path}")
+
+ ######################################################################
+ # Step 4: Deploy to Remote Device via RPC
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Connect to the remote device, upload the compiled library and parameters,
+ # then run inference remotely. This works for any device with TVM RPC server.
+ #
+ # Note: The following code demonstrates the RPC workflow. In local_demo mode,
+ # we skip actual execution to avoid LocalSession compatibility issues.
+
+ if local_demo:
+ # For demonstration, show the code structure without execution
+ print("\nRPC workflow (works for any remote device):")
+ print("=" * 50)
+ print("1. Start RPC server on target device:")
+ print(" python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090")
+ print("\n2. Connect from local machine:")
+ print(" remote = rpc.connect('DEVICE_IP', 9090)")
+ print("\n3. Upload compiled library:")
+ print(" remote.upload('model_deployed.so')")
+ print(" remote.upload('model_params.npz')")
+ print("\n4. Load and run remotely:")
+ print(" lib = remote.load_module('model_deployed.so')")
+ print(" vm = relax.VirtualMachine(lib, remote.cpu())")
+ print(" result = vm['main'](input, *params)")
+ print("\nDevice examples:")
+ print(" - Raspberry Pi: 192.168.1.100")
+ print(" - Remote server: ssh tunnel or direct IP")
+ print(" - NVIDIA Jetson: 10.0.0.50")
+ print(" - Cloud instance: public IP")
+ print("\nTo run actual RPC, set local_demo=False")
+ return # Skip actual RPC execution in demo mode
+
+ # Actual RPC workflow for real deployment
+ # Connect to remote device (works for ARM, x86, RISC-V, etc.)
+ # Make sure the RPC server is running on the device:
+ # python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090
+ device_host = "192.168.1.100" # Replace with your device IP
+ device_port = 9090
+ remote = rpc.connect(device_host, device_port)
+ print(f"Connected to remote device at {device_host}:{device_port}")
+
+ # Upload library and parameters to remote device
+ remote.upload(lib_path)
+ remote.upload(params_path)
+ print("Uploaded files to remote device")
+
+ # Load the library on the remote device
+ lib = remote.load_module("model_deployed.so")
+
+ # Choose device on remote machine
+ # For CPU: dev = remote.cpu()
+ # For CUDA GPU: dev = remote.cuda(0)
+ # For OpenCL: dev = remote.cl(0)
+ dev = remote.cpu()
+
+ # Create VM and load parameters
+ vm = relax.VirtualMachine(lib, dev)
+
+ # Load parameters from the uploaded file
+ # Note: In practice, you might load this from the remote filesystem
+ params_npz = np.load(params_path)
+ remote_params = [tvm.runtime.tensor(params_npz[f"p_{i}"], dev) for i in range(len(params_npz))]
+
+ ######################################################################
+ # Step 5: Run Inference on Remote Device
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Execute the model on the remote ARM device and retrieve results
+
+ # Prepare input data
+ input_data = np.random.randn(1, 1, 28, 28).astype("float32")
+ remote_input = tvm.runtime.tensor(input_data, dev)
+
+ # Run inference on remote device
+ output = vm["main"](remote_input, *remote_params)
+
+ # Extract result (handle both tuple and single tensor outputs)
+ if hasattr(output, "__len__") and len(output) > 0:
+ result = output[0]
+ else:
+ result = output
Review Comment:
The current check `hasattr(output, "__len__")` relies on duck-typing to distinguish between a tuple-like return and a single tensor. A more explicit and robust approach is to check `isinstance(output, tvm.runtime.Array)`, which is the specific type returned by the Relax VM for tuple outputs. This makes the code's intent clearer and is less prone to side effects.
```suggestion
if isinstance(output, tvm.runtime.Array) and len(output) > 0:
result = output[0]
else:
result = output
```
##########
docs/how_to/tutorials/cross_compilation_and_rpc.py:
##########
@@ -256,13 +256,328 @@ def run_opencl():
print("OpenCL test passed!")
+#########################################################################
+# Deploy PyTorch Models to Remote Devices with RPC
+# ------------------------------------------------
+# The above examples demonstrate cross compilation and RPC using low-level
+# TensorIR (via TE). For deploying complete neural network models from frameworks
+# like PyTorch or ONNX, TVM's Relax provides a higher-level abstraction that is
+# better suited for end-to-end model compilation.
+#
+# This section shows a modern workflow for deploying models to **any remote device**:
+#
+# 1. Import a PyTorch model and convert it to Relax
+# 2. Cross-compile for the target architecture (ARM, x86, RISC-V, etc.)
+# 3. Deploy via RPC to a remote device
+# 4. Run inference remotely
+#
+# This workflow is applicable to various deployment scenarios:
+#
+# - **ARM devices**: Raspberry Pi, NVIDIA Jetson, mobile phones
+# - **x86 servers**: Remote Linux servers, cloud instances
+# - **Embedded systems**: RISC-V boards, custom hardware
+# - **Accelerators**: Remote machines with GPUs, TPUs, or other accelerators
+#
+# .. note::
+# This example uses PyTorch for demonstration, but the workflow is identical
+# for ONNX models. Simply replace ``from_exported_program()`` with
+# ``from_onnx(model, keep_params_in_input=True)`` and follow the same steps.
+
+# First, let's check if PyTorch is available
+try:
+ import torch
+ from torch.export import export
+
+ HAS_TORCH = True
+except ImportError:
+ HAS_TORCH = False
+ print("PyTorch not available, skipping Relax example")
+
+
+def run_pytorch_model_via_rpc():
+ """
+ Demonstrates the complete workflow of deploying a PyTorch model to an ARM device via RPC.
+ """
+ if not HAS_TORCH:
+ print("Skipping PyTorch example (PyTorch not installed)")
+ return
+
+ from tvm import relax
+ from tvm.relax.frontend.torch import from_exported_program
+
+ ######################################################################
+ # Step 1: Define and Export PyTorch Model
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # We use a simple MLP model for demonstration. In practice, this could be
+ # any PyTorch model (ResNet, BERT, etc.).
+
+ class TorchMLP(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.net = torch.nn.Sequential(
+ torch.nn.Flatten(),
+ torch.nn.Linear(28 * 28, 128),
+ torch.nn.ReLU(),
+ torch.nn.Linear(128, 10),
+ )
+
+ def forward(self, data: torch.Tensor) -> torch.Tensor:
+ return self.net(data)
+
+ # Export the model using PyTorch 2.x export API
+ torch_model = TorchMLP().eval()
+ example_args = (torch.randn(1, 1, 28, 28, dtype=torch.float32),)
+
+ with torch.no_grad():
+ exported_program = export(torch_model, example_args)
+
+ ######################################################################
+ # Step 2: Convert to Relax and Prepare for Compilation
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Convert the exported PyTorch program to TVM's Relax representation
+
+ mod = from_exported_program(exported_program, keep_params_as_input=True)
+ # Separate parameters from the model for flexible deployment
+ mod, params = relax.frontend.detach_params(mod)
+
+ print("Converted PyTorch model to Relax:")
+ print(f" - Number of parameters: {len(params['main'])}")
+
+ ######################################################################
+ # Step 3: Cross-Compile for Target Device
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Compile the model for the target device architecture. The target
+ # configuration depends on your deployment scenario.
+
+ if local_demo:
+ # For demonstration on local machine, use local target
+ target = tvm.target.Target("llvm")
+ print("Using local target for demonstration")
+ else:
+ # Choose the appropriate target for your device:
+ #
+ # ARM devices:
+ # - Raspberry Pi 3/4 (32-bit): "llvm -mtriple=armv7l-linux-gnueabihf"
+ # - Raspberry Pi 4 (64-bit) / Jetson: "llvm -mtriple=aarch64-linux-gnu"
+ # - Android: "llvm -mtriple=aarch64-linux-android"
+ #
+ # x86 servers:
+ # - Linux x86_64: "llvm -mtriple=x86_64-linux-gnu"
+ # - With AVX-512: "llvm -mtriple=x86_64-linux-gnu -mcpu=skylake-avx512"
+ #
+ # RISC-V:
+ # - RV64: "llvm -mtriple=riscv64-unknown-linux-gnu"
+ #
+ # GPU targets:
+ # - CUDA: tvm.target.Target("cuda", host="llvm -mtriple=x86_64-linux-gnu")
+ # - OpenCL: tvm.target.Target("opencl", host="llvm -mtriple=aarch64-linux-gnu")
+ #
+ # For this example, we use ARM 64-bit
+ target = tvm.target.Target("llvm -mtriple=aarch64-linux-gnu")
+ print(f"Cross-compiling for target: {target}")
+
+ # Apply optimization pipeline
+ pipeline = relax.get_pipeline()
+ with target:
+ built_mod = pipeline(mod)
+
+ # Compile to executable
+ executable = tvm.compile(built_mod, target=target)
+
+ # Export to shared library
+ lib_path = temp.relpath("model_deployed.so")
+ executable.export_library(lib_path)
+ print(f"Exported library to: {lib_path}")
+
+ # Save parameters separately
+ import numpy as np
+
+ params_path = temp.relpath("model_params.npz")
+ param_arrays = {f"p_{i}": p.numpy() for i, p in enumerate(params["main"])}
+ np.savez(params_path, **param_arrays)
+ print(f"Saved parameters to: {params_path}")
+
+ ######################################################################
+ # Step 4: Deploy to Remote Device via RPC
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Connect to the remote device, upload the compiled library and parameters,
+ # then run inference remotely. This works for any device with TVM RPC server.
+ #
+ # Note: The following code demonstrates the RPC workflow. In local_demo mode,
+ # we skip actual execution to avoid LocalSession compatibility issues.
+
+ if local_demo:
+ # For demonstration, show the code structure without execution
+ print("\nRPC workflow (works for any remote device):")
+ print("=" * 50)
+ print("1. Start RPC server on target device:")
+ print(" python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090")
+ print("\n2. Connect from local machine:")
+ print(" remote = rpc.connect('DEVICE_IP', 9090)")
+ print("\n3. Upload compiled library:")
+ print(" remote.upload('model_deployed.so')")
+ print(" remote.upload('model_params.npz')")
+ print("\n4. Load and run remotely:")
+ print(" lib = remote.load_module('model_deployed.so')")
+ print(" vm = relax.VirtualMachine(lib, remote.cpu())")
+ print(" result = vm['main'](input, *params)")
+ print("\nDevice examples:")
+ print(" - Raspberry Pi: 192.168.1.100")
+ print(" - Remote server: ssh tunnel or direct IP")
+ print(" - NVIDIA Jetson: 10.0.0.50")
+ print(" - Cloud instance: public IP")
+ print("\nTo run actual RPC, set local_demo=False")
+ return # Skip actual RPC execution in demo mode
+
+ # Actual RPC workflow for real deployment
+ # Connect to remote device (works for ARM, x86, RISC-V, etc.)
+ # Make sure the RPC server is running on the device:
+ # python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090
+ device_host = "192.168.1.100" # Replace with your device IP
+ device_port = 9090
+ remote = rpc.connect(device_host, device_port)
+ print(f"Connected to remote device at {device_host}:{device_port}")
+
+ # Upload library and parameters to remote device
+ remote.upload(lib_path)
+ remote.upload(params_path)
+ print("Uploaded files to remote device")
+
+ # Load the library on the remote device
+ lib = remote.load_module("model_deployed.so")
+
+ # Choose device on remote machine
+ # For CPU: dev = remote.cpu()
+ # For CUDA GPU: dev = remote.cuda(0)
+ # For OpenCL: dev = remote.cl(0)
+ dev = remote.cpu()
+
+ # Create VM and load parameters
+ vm = relax.VirtualMachine(lib, dev)
+
+ # Load parameters from the uploaded file
+ # Note: In practice, you might load this from the remote filesystem
+ params_npz = np.load(params_path)
+ remote_params = [tvm.runtime.tensor(params_npz[f"p_{i}"], dev) for i in range(len(params_npz))]
+
+ ######################################################################
+ # Step 5: Run Inference on Remote Device
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Execute the model on the remote ARM device and retrieve results
+
+ # Prepare input data
+ input_data = np.random.randn(1, 1, 28, 28).astype("float32")
+ remote_input = tvm.runtime.tensor(input_data, dev)
+
+ # Run inference on remote device
+ output = vm["main"](remote_input, *remote_params)
+
+ # Extract result (handle both tuple and single tensor outputs)
+ if hasattr(output, "__len__") and len(output) > 0:
+ result = output[0]
+ else:
+ result = output
+
+ # Retrieve result from remote device to local
+ result_np = result.numpy()
+ print(f"Inference completed on remote device")
+ print(f" Output shape: {result_np.shape}")
+ print(f" Predicted class: {np.argmax(result_np)}")
+
+ ######################################################################
+ # Step 6: Performance Evaluation (Optional)
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Measure inference time on the remote device, excluding network overhead
+
+ time_f = vm.time_evaluator("main", dev, number=10, repeat=3)
+ prof_res = time_f(remote_input, *remote_params)
+ print(f"Inference time on remote device: {prof_res.mean * 1000:.2f} ms")
+
+ ######################################################################
+ # Notes on Performance Optimization
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ #
+ # For optimal performance on target devices, consider:
+ #
+ # 1. **Auto-tuning with MetaSchedule**: Use automated search to find
+ # optimal schedules for your specific hardware:
+ #
+ # .. code-block:: python
+ #
+ # mod = relax.get_pipeline(
+ # "static_shape_tuning",
+ # target=target,
+ # total_trials=2000
+ # )(mod)
+ #
+ # 2. **Quick optimization with DLight**: Apply pre-defined performant schedules:
+ #
+ # .. code-block:: python
+ #
+ # from tvm import dlight as dl
+ # with target:
+ # mod = dl.ApplyDefaultSchedule()(mod)
+ #
+ # 3. **Architecture-specific optimizations**:
+ #
+ # - ARM NEON SIMD: ``-mattr=+neon``
+ # - x86 AVX-512: ``-mcpu=skylake-avx512``
+ # - RISC-V Vector: ``-mattr=+v``
+ #
+ # .. code-block:: python
+ #
+ # # Example: ARM with NEON
+ # target = tvm.target.Target(
+ # "llvm -mtriple=aarch64-linux-gnu -mattr=+neon"
+ # )
+ #
+ # # Example: x86 with AVX-512
+ # target = tvm.target.Target(
+ # "llvm -mtriple=x86_64-linux-gnu -mcpu=skylake-avx512"
+ # )
+ #
+ # See :doc:`e2e_opt_model </how_to/tutorials/e2e_opt_model>` for detailed
+ # tuning examples.
+
+
+# Run the PyTorch RPC example if PyTorch is available
+if HAS_TORCH and local_demo:
Review Comment:
The condition `if HAS_TORCH and local_demo:` prevents the PyTorch RPC example from running when `local_demo` is set to `False`. However, users are instructed to set `local_demo` to `False` to test the actual RPC functionality. The `local_demo` logic is already handled within the `run_pytorch_model_via_rpc` function. To enable the example to run in both demo and real RPC modes, the condition should only check for `HAS_TORCH`.
```suggestion
if HAS_TORCH:
```
##########
docs/how_to/tutorials/cross_compilation_and_rpc.py:
##########
@@ -256,13 +256,328 @@ def run_opencl():
print("OpenCL test passed!")
+#########################################################################
+# Deploy PyTorch Models to Remote Devices with RPC
+# ------------------------------------------------
+# The above examples demonstrate cross compilation and RPC using low-level
+# TensorIR (via TE). For deploying complete neural network models from frameworks
+# like PyTorch or ONNX, TVM's Relax provides a higher-level abstraction that is
+# better suited for end-to-end model compilation.
+#
+# This section shows a modern workflow for deploying models to **any remote device**:
+#
+# 1. Import a PyTorch model and convert it to Relax
+# 2. Cross-compile for the target architecture (ARM, x86, RISC-V, etc.)
+# 3. Deploy via RPC to a remote device
+# 4. Run inference remotely
+#
+# This workflow is applicable to various deployment scenarios:
+#
+# - **ARM devices**: Raspberry Pi, NVIDIA Jetson, mobile phones
+# - **x86 servers**: Remote Linux servers, cloud instances
+# - **Embedded systems**: RISC-V boards, custom hardware
+# - **Accelerators**: Remote machines with GPUs, TPUs, or other accelerators
+#
+# .. note::
+# This example uses PyTorch for demonstration, but the workflow is identical
+# for ONNX models. Simply replace ``from_exported_program()`` with
+# ``from_onnx(model, keep_params_in_input=True)`` and follow the same steps.
+
+# First, let's check if PyTorch is available
+try:
+ import torch
+ from torch.export import export
+
+ HAS_TORCH = True
+except ImportError:
+ HAS_TORCH = False
+ print("PyTorch not available, skipping Relax example")
+
+
+def run_pytorch_model_via_rpc():
+ """
+ Demonstrates the complete workflow of deploying a PyTorch model to an ARM device via RPC.
+ """
+ if not HAS_TORCH:
+ print("Skipping PyTorch example (PyTorch not installed)")
+ return
+
+ from tvm import relax
+ from tvm.relax.frontend.torch import from_exported_program
+
+ ######################################################################
+ # Step 1: Define and Export PyTorch Model
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # We use a simple MLP model for demonstration. In practice, this could be
+ # any PyTorch model (ResNet, BERT, etc.).
+
+ class TorchMLP(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.net = torch.nn.Sequential(
+ torch.nn.Flatten(),
+ torch.nn.Linear(28 * 28, 128),
+ torch.nn.ReLU(),
+ torch.nn.Linear(128, 10),
+ )
+
+ def forward(self, data: torch.Tensor) -> torch.Tensor:
+ return self.net(data)
+
+ # Export the model using PyTorch 2.x export API
+ torch_model = TorchMLP().eval()
+ example_args = (torch.randn(1, 1, 28, 28, dtype=torch.float32),)
+
+ with torch.no_grad():
+ exported_program = export(torch_model, example_args)
+
+ ######################################################################
+ # Step 2: Convert to Relax and Prepare for Compilation
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Convert the exported PyTorch program to TVM's Relax representation
+
+ mod = from_exported_program(exported_program, keep_params_as_input=True)
+ # Separate parameters from the model for flexible deployment
+ mod, params = relax.frontend.detach_params(mod)
+
+ print("Converted PyTorch model to Relax:")
+ print(f" - Number of parameters: {len(params['main'])}")
+
+ ######################################################################
+ # Step 3: Cross-Compile for Target Device
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Compile the model for the target device architecture. The target
+ # configuration depends on your deployment scenario.
+
+ if local_demo:
+ # For demonstration on local machine, use local target
+ target = tvm.target.Target("llvm")
+ print("Using local target for demonstration")
+ else:
+ # Choose the appropriate target for your device:
+ #
+ # ARM devices:
+ # - Raspberry Pi 3/4 (32-bit): "llvm -mtriple=armv7l-linux-gnueabihf"
+ # - Raspberry Pi 4 (64-bit) / Jetson: "llvm -mtriple=aarch64-linux-gnu"
+ # - Android: "llvm -mtriple=aarch64-linux-android"
+ #
+ # x86 servers:
+ # - Linux x86_64: "llvm -mtriple=x86_64-linux-gnu"
+ # - With AVX-512: "llvm -mtriple=x86_64-linux-gnu -mcpu=skylake-avx512"
+ #
+ # RISC-V:
+ # - RV64: "llvm -mtriple=riscv64-unknown-linux-gnu"
+ #
+ # GPU targets:
+ # - CUDA: tvm.target.Target("cuda", host="llvm -mtriple=x86_64-linux-gnu")
+ # - OpenCL: tvm.target.Target("opencl", host="llvm -mtriple=aarch64-linux-gnu")
+ #
+ # For this example, we use ARM 64-bit
+ target = tvm.target.Target("llvm -mtriple=aarch64-linux-gnu")
+ print(f"Cross-compiling for target: {target}")
+
+ # Apply optimization pipeline
+ pipeline = relax.get_pipeline()
+ with target:
+ built_mod = pipeline(mod)
+
+ # Compile to executable
+ executable = tvm.compile(built_mod, target=target)
Review Comment:
While `tvm.compile` is a valid way to build the module, using `relax.build` is more idiomatic and specific for a Relax-centric workflow. This makes the tutorial clearer and more consistent with other Relax examples.
```suggestion
executable = relax.build(built_mod, target=target)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Docs][Relax] Update cross-compilation and RPC tutorial with modern PyTorch deployment workflow [tvm]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on PR #18413:
URL: https://github.com/apache/tvm/pull/18413#issuecomment-3476745945
## Summary of Changes
Hello @tlopex, I'm Gemini Code Assist[^1]! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request significantly enhances the `cross_compilation_and_rpc` tutorial by integrating a modern PyTorch model deployment workflow using TVM's Relax framework. It provides a high-level, end-to-end guide for users to import, compile, and deploy neural network models to diverse remote devices, complementing the existing low-level TensorIR examples. The update aims to make the tutorial more relevant and practical for contemporary model deployment scenarios.
### Highlights
* **New PyTorch/Relax Deployment Workflow**: Introduces a comprehensive guide for deploying PyTorch models to remote devices using TVM's Relax framework, covering model import, conversion, cross-compilation, and RPC deployment.
* **Modernized Tutorial Content**: Updates the existing cross-compilation and RPC tutorial to include a modern, high-level approach for end-to-end model compilation, complementing the existing low-level TensorIR examples.
* **Detailed Step-by-Step Guide**: Provides a clear, multi-step process for users, including defining and exporting PyTorch models, converting to Relax, cross-compiling for various architectures (ARM, x86, RISC-V), deploying via RPC, running remote inference, and evaluating performance.
* **Performance Optimization Guidance**: Includes practical recommendations for optimizing model performance on target devices, leveraging techniques such as MetaSchedule, DLight, and architecture-specific compiler flags.
<details>
<summary><b>Using Gemini Code Assist</b></summary>
<br>
The full guide for Gemini Code Assist can be found on our [documentation page](https://developers.google.com/gemini-code-assist/docs/review-github-code), here are some quick tips.
<b>Invoking Gemini</b>
You can request assistance from Gemini at any point by creating a comment using either `/gemini <command>` or `@gemini-code-assist <command>`. Below is a summary of the supported commands on the current page.
Feature | Command | Description
--- | --- | ---
Code Review | `/gemini review` | Performs a code review for the current pull request in its current state.
Pull Request Summary | `/gemini summary` | Provides a summary of the current pull request in its current state.
Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help | `/gemini help` | Displays a list of available commands.
<b>Customization</b>
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a `.gemini/` folder in the base of the repository. Detailed instructions can be found [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github).
<b>Limitations & Feedback</b>
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with :thumbsup: and :thumbsdown: on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up [here](https://google.qualtrics.com/jfe/form/SV_2cyuGuTWsEw84yG).
<b>You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the [Gemini Code Assist IDE Extension](https://cloud.google.com/products/gemini/code-assist).</b>
</details>
[^1]: Review the [Privacy Notices](https://policies.google.com/privacy), [Generative AI Prohibited Use Policy](https://policies.google.com/terms/generative-ai/use-policy), [Terms of Service](https://policies.google.com/terms), and learn how to configure Gemini Code Assist in GitHub [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). Gemini can make mistakes, so double check it and [use code with caution](https://support.google.com/legal/answer/13505487).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [DOCS] Update cross-compilation and RPC tutorial with modern PyTorch deployment workflow [tvm]
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on code in PR #18413:
URL: https://github.com/apache/tvm/pull/18413#discussion_r2483900615
##########
docs/how_to/tutorials/cross_compilation_and_rpc.py:
##########
@@ -256,13 +256,328 @@ def run_opencl():
print("OpenCL test passed!")
+#########################################################################
+# Deploy PyTorch Models to Remote Devices with RPC
+# ------------------------------------------------
+# The above examples demonstrate cross compilation and RPC using low-level
+# TensorIR (via TE). For deploying complete neural network models from frameworks
+# like PyTorch or ONNX, TVM's Relax provides a higher-level abstraction that is
+# better suited for end-to-end model compilation.
+#
+# This section shows a modern workflow for deploying models to **any remote device**:
+#
+# 1. Import a PyTorch model and convert it to Relax
+# 2. Cross-compile for the target architecture (ARM, x86, RISC-V, etc.)
+# 3. Deploy via RPC to a remote device
+# 4. Run inference remotely
+#
+# This workflow is applicable to various deployment scenarios:
+#
+# - **ARM devices**: Raspberry Pi, NVIDIA Jetson, mobile phones
+# - **x86 servers**: Remote Linux servers, cloud instances
+# - **Embedded systems**: RISC-V boards, custom hardware
+# - **Accelerators**: Remote machines with GPUs, TPUs, or other accelerators
+#
+# .. note::
+# This example uses PyTorch for demonstration, but the workflow is identical
+# for ONNX models. Simply replace ``from_exported_program()`` with
+# ``from_onnx(model, keep_params_in_input=True)`` and follow the same steps.
+
+# First, let's check if PyTorch is available
+try:
+ import torch
+ from torch.export import export
+
+ HAS_TORCH = True
+except ImportError:
+ HAS_TORCH = False
+ print("PyTorch not available, skipping Relax example")
+
+
+def run_pytorch_model_via_rpc():
+ """
+ Demonstrates the complete workflow of deploying a PyTorch model to an ARM device via RPC.
+ """
+ if not HAS_TORCH:
+ print("Skipping PyTorch example (PyTorch not installed)")
+ return
+
+ from tvm import relax
+ from tvm.relax.frontend.torch import from_exported_program
+
+ ######################################################################
+ # Step 1: Define and Export PyTorch Model
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # We use a simple MLP model for demonstration. In practice, this could be
+ # any PyTorch model (ResNet, BERT, etc.).
+
+ class TorchMLP(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.net = torch.nn.Sequential(
+ torch.nn.Flatten(),
+ torch.nn.Linear(28 * 28, 128),
+ torch.nn.ReLU(),
+ torch.nn.Linear(128, 10),
+ )
+
+ def forward(self, data: torch.Tensor) -> torch.Tensor:
+ return self.net(data)
+
+ # Export the model using PyTorch 2.x export API
+ torch_model = TorchMLP().eval()
+ example_args = (torch.randn(1, 1, 28, 28, dtype=torch.float32),)
+
+ with torch.no_grad():
+ exported_program = export(torch_model, example_args)
+
+ ######################################################################
+ # Step 2: Convert to Relax and Prepare for Compilation
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Convert the exported PyTorch program to TVM's Relax representation
+
+ mod = from_exported_program(exported_program, keep_params_as_input=True)
+ # Separate parameters from the model for flexible deployment
+ mod, params = relax.frontend.detach_params(mod)
+
+ print("Converted PyTorch model to Relax:")
+ print(f" - Number of parameters: {len(params['main'])}")
+
+ ######################################################################
+ # Step 3: Cross-Compile for Target Device
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Compile the model for the target device architecture. The target
+ # configuration depends on your deployment scenario.
+
+ if local_demo:
+ # For demonstration on local machine, use local target
+ target = tvm.target.Target("llvm")
+ print("Using local target for demonstration")
+ else:
+ # Choose the appropriate target for your device:
+ #
+ # ARM devices:
+ # - Raspberry Pi 3/4 (32-bit): "llvm -mtriple=armv7l-linux-gnueabihf"
+ # - Raspberry Pi 4 (64-bit) / Jetson: "llvm -mtriple=aarch64-linux-gnu"
+ # - Android: "llvm -mtriple=aarch64-linux-android"
+ #
+ # x86 servers:
+ # - Linux x86_64: "llvm -mtriple=x86_64-linux-gnu"
+ # - With AVX-512: "llvm -mtriple=x86_64-linux-gnu -mcpu=skylake-avx512"
+ #
+ # RISC-V:
+ # - RV64: "llvm -mtriple=riscv64-unknown-linux-gnu"
+ #
+ # GPU targets:
+ # - CUDA: tvm.target.Target("cuda", host="llvm -mtriple=x86_64-linux-gnu")
+ # - OpenCL: tvm.target.Target("opencl", host="llvm -mtriple=aarch64-linux-gnu")
+ #
+ # For this example, we use ARM 64-bit
+ target = tvm.target.Target("llvm -mtriple=aarch64-linux-gnu")
+ print(f"Cross-compiling for target: {target}")
+
+ # Apply optimization pipeline
+ pipeline = relax.get_pipeline()
+ with target:
+ built_mod = pipeline(mod)
+
+ # Compile to executable
+ executable = tvm.compile(built_mod, target=target)
+
+ # Export to shared library
+ lib_path = temp.relpath("model_deployed.so")
+ executable.export_library(lib_path)
+ print(f"Exported library to: {lib_path}")
+
+ # Save parameters separately
+ import numpy as np
+
+ params_path = temp.relpath("model_params.npz")
+ param_arrays = {f"p_{i}": p.numpy() for i, p in enumerate(params["main"])}
+ np.savez(params_path, **param_arrays)
+ print(f"Saved parameters to: {params_path}")
+
+ ######################################################################
+ # Step 4: Deploy to Remote Device via RPC
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Connect to the remote device, upload the compiled library and parameters,
+ # then run inference remotely. This works for any device with TVM RPC server.
+ #
+ # Note: The following code demonstrates the RPC workflow. In local_demo mode,
+ # we skip actual execution to avoid LocalSession compatibility issues.
+
+ if local_demo:
+ # For demonstration, show the code structure without execution
+ print("\nRPC workflow (works for any remote device):")
+ print("=" * 50)
+ print("1. Start RPC server on target device:")
+ print(" python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090")
+ print("\n2. Connect from local machine:")
+ print(" remote = rpc.connect('DEVICE_IP', 9090)")
+ print("\n3. Upload compiled library:")
+ print(" remote.upload('model_deployed.so')")
+ print(" remote.upload('model_params.npz')")
+ print("\n4. Load and run remotely:")
+ print(" lib = remote.load_module('model_deployed.so')")
+ print(" vm = relax.VirtualMachine(lib, remote.cpu())")
+ print(" result = vm['main'](input, *params)")
+ print("\nDevice examples:")
+ print(" - Raspberry Pi: 192.168.1.100")
+ print(" - Remote server: ssh tunnel or direct IP")
+ print(" - NVIDIA Jetson: 10.0.0.50")
+ print(" - Cloud instance: public IP")
+ print("\nTo run actual RPC, set local_demo=False")
+ return # Skip actual RPC execution in demo mode
+
+ # Actual RPC workflow for real deployment
+ # Connect to remote device (works for ARM, x86, RISC-V, etc.)
+ # Make sure the RPC server is running on the device:
+ # python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090
+ device_host = "192.168.1.100" # Replace with your device IP
+ device_port = 9090
+ remote = rpc.connect(device_host, device_port)
+ print(f"Connected to remote device at {device_host}:{device_port}")
+
+ # Upload library and parameters to remote device
+ remote.upload(lib_path)
+ remote.upload(params_path)
+ print("Uploaded files to remote device")
+
+ # Load the library on the remote device
+ lib = remote.load_module("model_deployed.so")
+
+ # Choose device on remote machine
+ # For CPU: dev = remote.cpu()
+ # For CUDA GPU: dev = remote.cuda(0)
+ # For OpenCL: dev = remote.cl(0)
+ dev = remote.cpu()
+
+ # Create VM and load parameters
+ vm = relax.VirtualMachine(lib, dev)
+
+ # Load parameters from the uploaded file
+ # Note: In practice, you might load this from the remote filesystem
+ params_npz = np.load(params_path)
+ remote_params = [tvm.runtime.tensor(params_npz[f"p_{i}"], dev) for i in range(len(params_npz))]
+
+ ######################################################################
+ # Step 5: Run Inference on Remote Device
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Execute the model on the remote ARM device and retrieve results
+
+ # Prepare input data
+ input_data = np.random.randn(1, 1, 28, 28).astype("float32")
+ remote_input = tvm.runtime.tensor(input_data, dev)
+
+ # Run inference on remote device
+ output = vm["main"](remote_input, *remote_params)
+
+ # Extract result (handle both tuple and single tensor outputs)
+ if hasattr(output, "__len__") and len(output) > 0:
+ result = output[0]
+ else:
+ result = output
Review Comment:
@tlopex this review comment is reasonable
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [DOCS] Update cross-compilation and RPC tutorial with modern PyTorch deployment workflow [tvm]
Posted by "tlopex (via GitHub)" <gi...@apache.org>.
tlopex commented on code in PR #18413:
URL: https://github.com/apache/tvm/pull/18413#discussion_r2483901170
##########
docs/how_to/tutorials/cross_compilation_and_rpc.py:
##########
@@ -256,13 +256,328 @@ def run_opencl():
print("OpenCL test passed!")
+#########################################################################
+# Deploy PyTorch Models to Remote Devices with RPC
+# ------------------------------------------------
+# The above examples demonstrate cross compilation and RPC using low-level
+# TensorIR (via TE). For deploying complete neural network models from frameworks
+# like PyTorch or ONNX, TVM's Relax provides a higher-level abstraction that is
+# better suited for end-to-end model compilation.
+#
+# This section shows a modern workflow for deploying models to **any remote device**:
+#
+# 1. Import a PyTorch model and convert it to Relax
+# 2. Cross-compile for the target architecture (ARM, x86, RISC-V, etc.)
+# 3. Deploy via RPC to a remote device
+# 4. Run inference remotely
+#
+# This workflow is applicable to various deployment scenarios:
+#
+# - **ARM devices**: Raspberry Pi, NVIDIA Jetson, mobile phones
+# - **x86 servers**: Remote Linux servers, cloud instances
+# - **Embedded systems**: RISC-V boards, custom hardware
+# - **Accelerators**: Remote machines with GPUs, TPUs, or other accelerators
+#
+# .. note::
+# This example uses PyTorch for demonstration, but the workflow is identical
+# for ONNX models. Simply replace ``from_exported_program()`` with
+# ``from_onnx(model, keep_params_in_input=True)`` and follow the same steps.
+
+# First, let's check if PyTorch is available
+try:
+ import torch
+ from torch.export import export
+
+ HAS_TORCH = True
+except ImportError:
+ HAS_TORCH = False
+ print("PyTorch not available, skipping Relax example")
+
+
+def run_pytorch_model_via_rpc():
+ """
+ Demonstrates the complete workflow of deploying a PyTorch model to an ARM device via RPC.
+ """
+ if not HAS_TORCH:
+ print("Skipping PyTorch example (PyTorch not installed)")
+ return
+
+ from tvm import relax
+ from tvm.relax.frontend.torch import from_exported_program
+
+ ######################################################################
+ # Step 1: Define and Export PyTorch Model
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # We use a simple MLP model for demonstration. In practice, this could be
+ # any PyTorch model (ResNet, BERT, etc.).
+
+ class TorchMLP(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.net = torch.nn.Sequential(
+ torch.nn.Flatten(),
+ torch.nn.Linear(28 * 28, 128),
+ torch.nn.ReLU(),
+ torch.nn.Linear(128, 10),
+ )
+
+ def forward(self, data: torch.Tensor) -> torch.Tensor:
+ return self.net(data)
+
+ # Export the model using PyTorch 2.x export API
+ torch_model = TorchMLP().eval()
+ example_args = (torch.randn(1, 1, 28, 28, dtype=torch.float32),)
+
+ with torch.no_grad():
+ exported_program = export(torch_model, example_args)
+
+ ######################################################################
+ # Step 2: Convert to Relax and Prepare for Compilation
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Convert the exported PyTorch program to TVM's Relax representation
+
+ mod = from_exported_program(exported_program, keep_params_as_input=True)
+ # Separate parameters from the model for flexible deployment
+ mod, params = relax.frontend.detach_params(mod)
+
+ print("Converted PyTorch model to Relax:")
+ print(f" - Number of parameters: {len(params['main'])}")
+
+ ######################################################################
+ # Step 3: Cross-Compile for Target Device
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Compile the model for the target device architecture. The target
+ # configuration depends on your deployment scenario.
+
+ if local_demo:
+ # For demonstration on local machine, use local target
+ target = tvm.target.Target("llvm")
+ print("Using local target for demonstration")
+ else:
+ # Choose the appropriate target for your device:
+ #
+ # ARM devices:
+ # - Raspberry Pi 3/4 (32-bit): "llvm -mtriple=armv7l-linux-gnueabihf"
+ # - Raspberry Pi 4 (64-bit) / Jetson: "llvm -mtriple=aarch64-linux-gnu"
+ # - Android: "llvm -mtriple=aarch64-linux-android"
+ #
+ # x86 servers:
+ # - Linux x86_64: "llvm -mtriple=x86_64-linux-gnu"
+ # - With AVX-512: "llvm -mtriple=x86_64-linux-gnu -mcpu=skylake-avx512"
+ #
+ # RISC-V:
+ # - RV64: "llvm -mtriple=riscv64-unknown-linux-gnu"
+ #
+ # GPU targets:
+ # - CUDA: tvm.target.Target("cuda", host="llvm -mtriple=x86_64-linux-gnu")
+ # - OpenCL: tvm.target.Target("opencl", host="llvm -mtriple=aarch64-linux-gnu")
+ #
+ # For this example, we use ARM 64-bit
+ target = tvm.target.Target("llvm -mtriple=aarch64-linux-gnu")
+ print(f"Cross-compiling for target: {target}")
+
+ # Apply optimization pipeline
+ pipeline = relax.get_pipeline()
+ with target:
+ built_mod = pipeline(mod)
+
+ # Compile to executable
+ executable = tvm.compile(built_mod, target=target)
+
+ # Export to shared library
+ lib_path = temp.relpath("model_deployed.so")
+ executable.export_library(lib_path)
+ print(f"Exported library to: {lib_path}")
+
+ # Save parameters separately
+ import numpy as np
+
+ params_path = temp.relpath("model_params.npz")
+ param_arrays = {f"p_{i}": p.numpy() for i, p in enumerate(params["main"])}
+ np.savez(params_path, **param_arrays)
+ print(f"Saved parameters to: {params_path}")
+
+ ######################################################################
+ # Step 4: Deploy to Remote Device via RPC
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Connect to the remote device, upload the compiled library and parameters,
+ # then run inference remotely. This works for any device with TVM RPC server.
+ #
+ # Note: The following code demonstrates the RPC workflow. In local_demo mode,
+ # we skip actual execution to avoid LocalSession compatibility issues.
+
+ if local_demo:
+ # For demonstration, show the code structure without execution
+ print("\nRPC workflow (works for any remote device):")
+ print("=" * 50)
+ print("1. Start RPC server on target device:")
+ print(" python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090")
+ print("\n2. Connect from local machine:")
+ print(" remote = rpc.connect('DEVICE_IP', 9090)")
+ print("\n3. Upload compiled library:")
+ print(" remote.upload('model_deployed.so')")
+ print(" remote.upload('model_params.npz')")
+ print("\n4. Load and run remotely:")
+ print(" lib = remote.load_module('model_deployed.so')")
+ print(" vm = relax.VirtualMachine(lib, remote.cpu())")
+ print(" result = vm['main'](input, *params)")
+ print("\nDevice examples:")
+ print(" - Raspberry Pi: 192.168.1.100")
+ print(" - Remote server: ssh tunnel or direct IP")
+ print(" - NVIDIA Jetson: 10.0.0.50")
+ print(" - Cloud instance: public IP")
+ print("\nTo run actual RPC, set local_demo=False")
+ return # Skip actual RPC execution in demo mode
+
+ # Actual RPC workflow for real deployment
+ # Connect to remote device (works for ARM, x86, RISC-V, etc.)
+ # Make sure the RPC server is running on the device:
+ # python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090
+ device_host = "192.168.1.100" # Replace with your device IP
+ device_port = 9090
+ remote = rpc.connect(device_host, device_port)
+ print(f"Connected to remote device at {device_host}:{device_port}")
+
+ # Upload library and parameters to remote device
+ remote.upload(lib_path)
+ remote.upload(params_path)
+ print("Uploaded files to remote device")
+
+ # Load the library on the remote device
+ lib = remote.load_module("model_deployed.so")
+
+ # Choose device on remote machine
+ # For CPU: dev = remote.cpu()
+ # For CUDA GPU: dev = remote.cuda(0)
+ # For OpenCL: dev = remote.cl(0)
+ dev = remote.cpu()
+
+ # Create VM and load parameters
+ vm = relax.VirtualMachine(lib, dev)
+
+ # Load parameters from the uploaded file
+ # Note: In practice, you might load this from the remote filesystem
+ params_npz = np.load(params_path)
+ remote_params = [tvm.runtime.tensor(params_npz[f"p_{i}"], dev) for i in range(len(params_npz))]
+
+ ######################################################################
+ # Step 5: Run Inference on Remote Device
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ # Execute the model on the remote ARM device and retrieve results
+
+ # Prepare input data
+ input_data = np.random.randn(1, 1, 28, 28).astype("float32")
+ remote_input = tvm.runtime.tensor(input_data, dev)
+
+ # Run inference on remote device
+ output = vm["main"](remote_input, *remote_params)
+
+ # Extract result (handle both tuple and single tensor outputs)
+ if hasattr(output, "__len__") and len(output) > 0:
+ result = output[0]
+ else:
+ result = output
Review Comment:
Okay
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org