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.