Edge AI Deployment

Optimization techniques to make models fit on edge devices - quantization, pruning, knowledge distillation - plus practical deployment to Raspberry Pi and mobile.

Model Optimization Techniques

Pruning

Pruning removes redundant parameters from a neural network. Many weights are near zero and contribute little to the output:

import torch
import torch.nn.utils.prune as prune

model = load_trained_model()

# Prune 50% of weights in each Conv2d layer
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Conv2d):
        prune.l1_unstructured(module, name='weight', amount=0.5)

# Check sparsity
total, zeros = 0, 0
for name, param in model.named_parameters():
    if 'weight' in name:
        total += param.numel()
        zeros += (param == 0).sum().item()
print(f"Sparsity: {zeros/total*100:.1f}%")

# Fine-tune the pruned model to recover accuracy
# Then make pruning permanent:
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Conv2d):
        prune.remove(module, 'weight')

Knowledge Distillation

Knowledge distillation trains a small "student" model to mimic a large "teacher" model. The student learns from the teacher's soft probability outputs, which contain more information than hard labels:

• Teacher: A large, accurate model (e.g., ResNet-152).
• Student: A small, fast model (e.g., MobileNet-v2).
• Temperature: Softmax temperature >1 makes the teacher's outputs smoother, revealing inter-class relationships.
• Typical accuracy recovery: The student achieves 95-99% of the teacher's accuracy at a fraction of the size.

Deploying to Raspberry Pi

1. Prepare the Model

   Convert and quantize your model to TFLite INT8 format on your development machine.

2. Set Up the Pi

   Install tflite-runtime: pip install tflite-runtime. No need for full TensorFlow.

3. Connect Camera

   Attach the Pi Camera Module. Use picamera2 library for capture.

4. Run Inference Loop

   Capture frames, preprocess, run TFLite inference, and display or act on results.

Mobile Deployment

Efficient Architectures for Edge

• MobileNet v2/v3: Inverted residuals and depthwise separable convolutions. The go-to architecture for mobile vision.
• EfficientNet-Lite: EfficientNet optimized for mobile. Removes squeeze-and-excitation blocks for better TFLite compatibility.
• YOLOv8-nano: Ultralytics' smallest YOLO model. Real-time object detection on mobile devices.
• DistilBERT: 60% smaller, 60% faster than BERT with 97% of its language understanding performance.