Guides

Environment Setup

Prerequisites

Before starting, ensure your system meets these requirements:

  • Linux operating system (Ubuntu 20.04 or later recommended)
  • Apropriate NVIDIA GPU
  • Apropriate NVIDIA driver
  • At least 10GB of free disk space

Table of Contents

  1. Install NVIDIA Driver
  2. Install Docker
  3. Install NVIDIA Container Toolkit
  4. Install TensorRT Container
  5. Verification
  6. Troubleshooting

Install NVIDIA Driver

  1. Check if NVIDIA driver is installed: 
    nvidia-smi
  2. Install NVIDIA drivers: 
    sudo apt install nvidia-driver-###  # Use latest recommended version
sudo reboot

Install Docker

  1. Check if docker is install and remove old versions (if any): 
    dpkg -l | grep -i docker
  2. Install Docker Engine: 
    sudo apt-get update
sudo apt-get install docker-ce docker-ce-cli containerd.io docker-compose-plugin
  3. Verify Docker installation: 
    sudo docker run hello-world

Install NVIDIA Container Toolkit

  1. Install NVIDIA Container Toolkit (if necessary):

Install TensorRT Container

  1. Pull the TensorRT container with TensorFlow support: 
    sudo docker pull nvcr.io/nvidia/tensorflow:##.##-tf#-py3
  2. Run the container: 
    sudo docker run --gpus all -it --rm nvcr.io/nvidia/tensorflow:##.##-tf#-py3

Verification

  1. Inside the container, verify TensorFlow can see the GPU: 
    import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
  2. Verify TensorRT: 
    import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt
print(tf.__version__)

Troubleshooting

Common Issues and Solutions

  1. Docker permission denied

  2. NVIDIA driver not found 
    sudo reboot
  3. Container fails to start with GPU error 
    # Check if NVIDIA runtime is properly configured
sudo docker info | grep -i runtime

Additional Resources

Note: Version numbers in this guide may need to be updated based on your system requirements and the latest available versions.

  1. Log-in into a Lambda server using your credentials via ssh.

  2. Once in the Lambda server execute the following command:

    jupyter notebook --no-browser --port=8080

    • This will open the lambda server and will provide a token.

  3. On your local systems type: ssh -L 8080:localhost:8080 <your account>@<lambda server>

    • This command is basically forwarding port 8080 on the lambda server to your local system.
  4. On your system browser go to: https://localhost:8080 and that will access the jupyter notebook interface from the lambda server.

Troubleshooting:

  • Port Usage: If there are other person using that port (e.g. 8080) the step 2 will fail. You need to select another port. If this happens, the new port must be used on step 4.

Hardware Setup

  1. First we ned to update the Firmware on the Jetson Orin Nano. If you are not sure which one you have. Recommend the follow step 2 to upgrade firmware to 36.x or greater. Otherwise, skip to step 4.
  2. In the following link for the Jetson Orin Nano Jetpack 5.13 firmware from step 2.1 (You can follow all the steps from this link and skip to step 5).
  3. Download the 5.13 Jetpack and follow step 2 in the dowload for your operating system. Make sure to flash with the Jetpack 5.13 firmware and not Jetpack 6.x.
    • You need to download Balena Etcher in the guide as well. This flashes the SD card with the firmware. You may need other software depending on your operating system.
  4. Then follow all the steps in write image to micro SD card. Make sure to download latest Jetpack (currently v6.1).
  5. Once the install Jetpack on top of the Jetson Linux: 
    apt install nvidia-jetpack

Model Conversion

PyTorch to LiteRT

Table of Contents

The litert_torch Version

This notebook converts a PyTorch model directly into a LiteRT .tflite file.

  1. Load the PyTorch model
  2. Convert it with litert_torch
  3. Export the generated .tflite file

NOTE: This code is also available in the repository, specifically notebooks/model_conversion/pytorch_to_litert.ipynb.

Setup

Install the dependencies in your python environment first, it is recommended to use a separate virtual environment for this notebook compared to the one used to install the dependencies of requirements.txt since some of these have specific version requirements that may conflict with the ones found there.

Especifically since litert_torch requires torch==2.6.0 or greater, which differs with the version required by requirements.txt.

%pip install --upgrade torch torchvision litert_torch jupyter ipykernel

Run to check that the dependencies are correctly installed.

from pathlib import Path

import torch

from torchvision.models import mobilenet_v2
from ai_edge_litert import interpreter as litert_interpreter
import litert_torch

print("Active PyTorch version:", torch.__version__)

ARTIFACTS_DIR = Path('artifacts')
ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)

WEIGHTS_PATH = ARTIFACTS_DIR / 'mobilenet_v2_best_model.pth'
TFLITE_PATH = ARTIFACTS_DIR / 'mobilenet_v2_litert.tflite'

print('Artifacts dir:', ARTIFACTS_DIR)

This artifacts directory is where the generated .tflite file will be saved, as well as where you can put your PyTorch model weights to be converted.

1) Load PyTorch model

Here we load the PyTorch model and extract the number of classes from the final layer, in order to initialize the LiteRT model with the correct output dimensions.

checkpoint = torch.load(WEIGHTS_PATH, map_location='cpu')

# In case the checkpoint is a dict with a 'state_dict' key
state_dict = (
    checkpoint['state_dict'] 
    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint 
    else checkpoint
)

# Determine the number of classes from the last tensor in the state dict 
last_tensor = state_dict[list(state_dict.keys())[-1]]

num_classes = (
    last_tensor.shape[0] 
    if last_tensor.ndim == 1        # bias vector case
    else min(last_tensor.shape)     # weight matrix case
)

model = mobilenet_v2(num_classes=num_classes)
model.load_state_dict(state_dict, strict=True)
model.eval()

print(model)

The example above assumes that the model being loaded is a MobileNetV2 architecture, but you can replace it with any other architecture.

2) Convert PyTorch -> LiteRT

Defining some sample inputs to trace the model and later convert it to LiteRT format with litert_torch and finally exporting the generated .tflite file to the artifacts directory.

sample_inputs = (torch.randn(1, 3, 224, 224),) # Must be a tuple

edge_model = litert_torch.convert(model, sample_inputs)
edge_model.export(str(TFLITE_PATH))

3) Final check

Just checking that the file was generated and inspecting its size.

print('TFLite path:', TFLITE_PATH)
print('TFLite size (bytes):', TFLITE_PATH.stat().st_size)

interpreter = litert_interpreter.Interpreter(model_path=str(TFLITE_PATH))
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

print("Input details:", input_details)
print("Output details:", output_details)

The ONNX Version (Deprecated)

This tutorial shows how to convert a PyTorch model into LiteRT format using ONNX as an intermediate format. It’s important to note that even though both of this methods work, the litert_torch method is the current recommended approach, since this one uses deprecated dependencies and could lead to compatibility issues in the future.

The conversion process involves three main steps:

  1. PyTorch model -> ONNX
  2. ONNX -> TensorFlow SavedModel
  3. SavedModel -> TFLite (.tflite)

Setup

Make sure you have the following dependencies installed in your Python environment, again it is recommended to use a separate virtual environment for this notebook compared to the one used to install the dependencies of requirements.txt since some of these have specific version requirements that may conflict with the ones found there.

torch>=2.5.1
torchvision==0.20.1
onnx==1.15.0
onnx-tf==1.10.0
tensorflow==2.15.1
tensorflow-probability==0.23.0

You can save these in a onnx_conversion.txt file and install them with.

pip install -r onnx_conversion.txt

Run to check that the dependencies are correctly installed.

from pathlib import Path

import torch

from torchvision.models import mobilenet_v2
import tensorflow as tf
import onnx
from onnx_tf.backend import prepare

ARTIFACTS_DIR = Path('artifacts')
ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)

ONNX_PATH = ARTIFACTS_DIR / 'model.onnx'
SAVEDMODEL_DIR = ARTIFACTS_DIR / 'saved_model'
TFLITE_PATH = ARTIFACTS_DIR / 'model.tflite'

WEIGHTS_PATH = ARTIFACTS_DIR / 'mobilenet_v2_best_model.pth'
TFLITE_PATH = ARTIFACTS_DIR / 'mobilenet_v2_litert.tflite'

print('Artifacts dir:', ARTIFACTS_DIR)

1) Load PyTorch model

Here we load the PyTorch model and extract the number of classes from the final layer, in order to initialize the ONNX model with the correct output dimensions.

checkpoint = torch.load(WEIGHTS_PATH, map_location="cpu")

# In case the checkpoint is a dict with a 'state_dict' key
state_dict = (
    checkpoint['state_dict'] 
    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint 
    else checkpoint
)

# Determine the number of classes from the last tensor in the state dict 
last_tensor = state_dict[list(state_dict.keys())[-1]]

num_classes = (
    last_tensor.shape[0] 
    if last_tensor.ndim == 1        # bias vector case
    else min(last_tensor.shape)     # weight matrix case
)

model = mobilenet_v2(num_classes=num_classes)
model.load_state_dict(state_dict, strict=True)
model.eval()

print(model)

2) Export PyTorch -> ONNX

Here we define some sample inputs to trace the model and then export it to ONNX format. The parameters used in torch.onnx.export are set to be compatible with ONNX Runtime and TensorFlow.

sample_inputs = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    sample_inputs,
    str(ONNX_PATH),
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
    dynamo=False,
)

print("ONNX exported to:", ONNX_PATH)

3) Convert ONNX -> TensorFlow SavedModel

This step uses onnx-tf to convert the ONNX model into a TensorFlow SavedModel, which is like a directory containing the model architecture and weights in a format that can be easily converted to TFLite.

onnx_model = onnx.load(str(ONNX_PATH))
tf_rep = prepare(onnx_model)
tf_rep.export_graph(str(SAVEDMODEL_DIR))

print("SavedModel exported to:", SAVEDMODEL_DIR)

4) Convert SavedModel -> LiteRT

Finally we use TensorFlow’s TFLiteConverter to convert the SavedModel into a .tflite file.

converter = tf.lite.TFLiteConverter.from_saved_model(str(SAVEDMODEL_DIR))

tflite_model = converter.convert()
TFLITE_PATH.write_bytes(tflite_model)

5) Final check

Just checking that the file was generated, as well as loading it with the TFLite interpreter to check that it’s valid.

print('TFLite path:', TFLITE_PATH)
print('TFLite size (bytes):', TFLITE_PATH.stat().st_size)

interpreter = tf.lite.Interpreter(model_path=str(TFLITE_PATH))
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

print("Input details:", input_details)
print("Output details:", output_details)
Guides