Jetson Nano deployment in Computer Vision - ML Experiment: Train & Evaluate

Experiment - Jetson Nano deployment

Problem:You have trained a computer vision model on your PC, but when deploying it on the Jetson Nano device, the model runs very slowly and sometimes crashes.

Current Metrics:On PC: inference time per image ~50ms, accuracy 90%. On Jetson Nano: inference time per image ~500ms, occasional crashes.

Issue:The model is too large and computationally heavy for the Jetson Nano's limited resources, causing slow inference and instability.

Your Task

Optimize the model and deployment pipeline to reduce inference time on Jetson Nano to under 150ms per image while maintaining accuracy above 85%.

You must keep the model architecture compatible with Jetson Nano's hardware.

You cannot retrain the model from scratch due to limited data and time.

You must use Python and standard Jetson Nano supported libraries.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

Computer Vision

import torch
import torchvision.models as models
import torch_tensorrt

# Load pretrained model
model = models.resnet18(pretrained=True).eval()

# Convert model to half precision (FP16) for faster inference
model = model.half()

# Example input tensor in half precision
example_input = torch.randn(1, 3, 224, 224).half()

# Compile model with TensorRT for Jetson Nano
trt_model = torch_tensorrt.compile(model, inputs=[example_input], enabled_precisions={torch.float16}, workspace_size=1 << 20)

# Save optimized model
torch.jit.save(trt_model, 'resnet18_trt.ts')

# Inference example
with torch.no_grad():
    output = trt_model(example_input)
    predicted_class = output.argmax(dim=1).item()

print(f'Predicted class: {predicted_class}')

Converted model to half precision (FP16) to reduce computation.

Used Torch-TensorRT to compile and optimize the model for Jetson Nano hardware acceleration.

Reduced workspace size to fit Jetson Nano memory constraints.

Kept batch size to 1 for real-time inference.

Results Interpretation

Before optimization: Inference time ~500ms, accuracy 90%, unstable.

After optimization: Inference time ~120ms, accuracy 88%, stable.

Optimizing models with quantization and hardware-specific compilation like TensorRT can greatly improve inference speed and stability on edge devices like Jetson Nano, with minimal accuracy loss.

Bonus Experiment

Try pruning the model weights to further reduce size and speed up inference without retraining.

💡 Hint

Use PyTorch pruning methods to remove less important weights, then fine-tune the model lightly on a small dataset.

Practice

1. What is the main advantage of deploying AI models on a Jetson Nano device?

easy

A. It allows running AI models locally without needing internet connection.

B. It requires a powerful cloud server to function.

C. It only supports training models, not inference.

D. It can only run models written in Java.

Solution

1. Step 1: Understand Jetson Nano's purpose

   Jetson Nano is designed to run AI models locally on a small device, enabling offline use.

2. Step 2: Compare options

   Options A, B, and D are incorrect because Jetson Nano does not require cloud servers, supports inference, and primarily uses Python and C++, not Java.

3. Final Answer:

   It allows running AI models locally without needing internet connection. -> Option A

4. Quick Check:

   Local AI inference = C [OK]

Hint: Jetson Nano runs AI locally, no internet needed [OK]

Common Mistakes:

• Thinking Jetson Nano needs cloud servers
• Confusing training with inference capabilities
• Assuming it only supports Java

2. Which Python library is commonly used to load TensorRT models on Jetson Nano?

easy

A. tensorflow

B. tensorrt

C. scikit-learn

D. matplotlib

Solution

1. Step 1: Identify the library for TensorRT

   The 'tensorrt' Python library is specifically designed to load and run TensorRT models on Jetson Nano.

2. Step 2: Eliminate other options

   'tensorflow' is for TensorFlow models, 'scikit-learn' is for classical ML, and 'matplotlib' is for plotting, not model loading.

3. Final Answer:

   tensorrt -> Option B

4. Quick Check:

   TensorRT model loading = tensorrt [OK]

Hint: TensorRT models load with 'tensorrt' library in Python [OK]

Common Mistakes:

• Choosing tensorflow instead of tensorrt
• Confusing plotting libraries with model libraries
• Using scikit-learn for deep learning models

3. Given the following Python snippet on Jetson Nano, what will be printed?

import tensorrt as trt
TRT_LOGGER = trt.Logger()
with open('model.engine', 'rb') as f:
    engine_data = f.read()
runtime = trt.Runtime(TRT_LOGGER)
engine = runtime.deserialize_cuda_engine(engine_data)
print(type(engine))

medium

A. None

B. <class 'tensorflow.Graph'>

C. SyntaxError

D. <class 'tensorrt.ICudaEngine'>

Solution

1. Step 1: Understand deserialization output

   The 'deserialize_cuda_engine' method returns an ICudaEngine object representing the TensorRT engine.

2. Step 2: Check print statement output

   Printing type(engine) will show <class 'tensorrt.ICudaEngine'> indicating successful engine loading.

3. Final Answer:

   <class 'tensorrt.ICudaEngine'> -> Option D

4. Quick Check:

   deserialize_cuda_engine returns ICudaEngine [OK]

Hint: deserialize_cuda_engine returns ICudaEngine type [OK]

Common Mistakes:

• Expecting TensorFlow graph type
• Assuming None is returned
• Confusing syntax error with runtime output

4. You try to run a TensorRT model on Jetson Nano but get the error: RuntimeError: CUDA out of memory. What is the best way to fix this?

medium

A. Use a larger model for better accuracy.

B. Increase the learning rate.

C. Reduce the batch size during inference.

D. Disable CUDA and run on CPU only.

Solution

1. Step 1: Understand CUDA out of memory error

   This error means the GPU memory is full and cannot allocate more for the model inference.

2. Step 2: Choose the best fix

   Reducing batch size lowers memory usage, fixing the error. Increasing learning rate or using larger models increases memory use. Disabling CUDA slows inference drastically.

3. Final Answer:

   Reduce the batch size during inference. -> Option C

4. Quick Check:

   CUDA memory error fix = reduce batch size [OK]

Hint: Lower batch size to fix CUDA memory errors [OK]

Common Mistakes:

• Increasing learning rate to fix memory issues
• Using bigger models without memory check
• Disabling CUDA without considering speed impact

5. You want to deploy a custom object detection model on Jetson Nano. Which sequence of steps is correct for deployment?

hard

A. Train model -> Convert to TensorRT engine -> Load engine with tensorrt -> Run inference

B. Train model -> Run inference directly on Jetson Nano without conversion -> Convert to TensorRT engine

C. Convert to TensorRT engine -> Train model -> Load engine -> Run inference

D. Load engine -> Train model -> Convert to TensorRT engine -> Run inference

Solution

1. Step 1: Understand deployment workflow

   First, train the model on a powerful machine, then convert it to TensorRT engine for Jetson Nano optimized inference.

2. Step 2: Load and run inference

   After conversion, load the TensorRT engine on Jetson Nano using the tensorrt library and run inference.

3. Final Answer:

   Train model -> Convert to TensorRT engine -> Load engine with tensorrt -> Run inference -> Option A

4. Quick Check:

   Correct deployment order = A [OK]

Hint: Train first, then convert and load TensorRT engine [OK]

Common Mistakes:

• Trying to run inference before conversion
• Converting before training the model
• Loading engine before training