YOLOv5 for Real-Time Inference: From Training to Edge Deployment

In nine years of building computer vision systems, the most challenging deployments have not been the ones with the most complex models—they have been the ones where the model had to run at the edge. A defect detection system on a factory floor running on a Jetson Nano. A wildlife monitoring camera in a remote forest preserve with intermittent connectivity. A retail analytics system running on an Intel NUC in a store backroom. In every case, the constraint was the same: the model had to be fast, small, and accurate enough to be useful, running on hardware that would make a data center GPU laugh.

YOLOv5 has been my weapon of choice for these deployments. Not because it is the absolute state-of-the-art in accuracy—newer architectures like YOLOv8 and YOLOv9 have pushed mAP higher on benchmarks—but because YOLOv5 offers the best combination of accuracy, speed, deployment tooling, and community support for production edge deployment. This is a practical guide to taking YOLOv5 from training to production on edge devices.

Why YOLOv5 for Edge Deployment

The YOLO family (You Only Look Once) revolutionized object detection by reframing it as a single regression problem: predict bounding boxes and class probabilities directly from the image in a single forward pass. YOLOv5, developed by Ultralytics, refined this approach with:

1. PyTorch-native implementation: Clean, readable code that is easy to modify, debug, and extend.
2. Auto-learning anchors: Automatic anchor box calculation based on your dataset, eliminating manual tuning.
3. Built-in export pipeline: First-class support for ONNX, CoreML, TFLite, TensorRT, and OpenVINO export.
4. Scalable architecture: From YOLOv5n (1.9M parameters) to YOLOv5x (86.7M parameters), you can choose the right accuracy/speed tradeoff for your hardware.
5. Active community and ecosystem: Extensive documentation, pretrained models, and community contributions.

Training YOLOv5

Dataset Preparation

YOLOv5 expects a specific directory structure and annotation format:

dataset/
├── train/
│   ├── images/
│   │   ├── img001.jpg
│   │   └── img002.jpg
│   └── labels/
│       ├── img001.txt
│       └── img002.txt
├── val/
│   ├── images/
│   └── labels/
└── data.yaml

Each label file contains one line per object:

# class_id center_x center_y width height (normalized 0-1)
0 0.45 0.52 0.12 0.08
1 0.78 0.33 0.05 0.15

The data.yaml configuration:

path: /data/my_dataset
train: train/images
val: val/images
test: test/images

nc: 5
names: ['crack', 'corrosion', 'dent', 'scratch', 'deformation']

Converting from Other Formats

Most datasets are not in YOLO format. Here is a reliable conversion pipeline:

import json
import cv2
import os

def coco_to_yolo(coco_json_path, output_dir):
    """Convert COCO annotations to YOLO format."""
    with open(coco_json_path) as f:
        coco = json.load(f)

    # Build category mapping
    cat_map = {cat["id"]: idx for idx, cat in enumerate(coco["categories"])}

    # Build image id to filename mapping
    img_map = {img["id"]: img for img in coco["images"]}

    # Group annotations by image
    from collections import defaultdict
    anns_by_img = defaultdict(list)
    for ann in coco["annotations"]:
        anns_by_img[ann["image_id"]].append(ann)

    os.makedirs(f"{output_dir}/labels", exist_ok=True)

    for img_id, img_info in img_map.items():
        w, h = img_info["width"], img_info["height"]
        filename = os.path.splitext(img_info["file_name"])[0]

        lines = []
        for ann in anns_by_img.get(img_id, []):
            cls_id = cat_map[ann["category_id"]]
            # COCO bbox: [x_top_left, y_top_left, width, height]
            # YOLO bbox: [center_x, center_y, width, height] (normalized)
            x, y, bw, bh = ann["bbox"]
            cx = (x + bw / 2) / w
            cy = (y + bh / 2) / h
            nw = bw / w
            nh = bh / h
            lines.append(f"{cls_id} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}")

        with open(f"{output_dir}/labels/{filename}.txt", "w") as f:
            f.write("\n".join(lines))

Training Configuration

# Train YOLOv5s (small) for edge deployment
python train.py \
    --img 640 \
    --batch 32 \
    --epochs 300 \
    --data data.yaml \
    --weights yolov5s.pt \
    --cfg yolov5s.yaml \
    --name defect_detection \
    --hyp data/hyps/hyp.scratch-low.yaml \
    --patience 50 \
    --save-period 50 \
    --device 0

Hyperparameter Tuning for Edge Models

For edge deployment, I use specific hyperparameter strategies:

# hyp.edge.yaml - Optimized for edge deployment
lr0: 0.01
lrf: 0.01
momentum: 0.937
weight_decay: 0.0005
warmup_epochs: 3.0
warmup_momentum: 0.8
warmup_bias_lr: 0.1
box: 0.05
cls: 0.5
cls_pw: 1.0
obj: 1.0
obj_pw: 1.0
iou_t: 0.20
anchor_t: 4.0
fl_gamma: 0.0
hsv_h: 0.015
hsv_s: 0.7
hsv_v: 0.4
degrees: 0.0
translate: 0.1
scale: 0.5
shear: 0.0
perspective: 0.0
flipud: 0.0
fliplr: 0.5
mosaic: 1.0
mixup: 0.0
copy_paste: 0.0

Key decisions:

• Lower learning rate (0.01 vs default 0.01): More stable convergence for smaller models.
• More augmentation: scale: 0.5 and translate: 0.1 help the smaller model generalize better.
• No copy-paste or mixup: These augmentations can hurt small models that lack capacity.

Transfer Learning Strategy

# Stage 1: Train head only (freeze backbone)
python train.py \
    --weights yolov5s.pt \
    --freeze 10 \
    --epochs 50 \
    --data data.yaml

# Stage 2: Fine-tune entire network
python train.py \
    --weights runs/train/exp/weights/last.pt \
    --epochs 200 \
    --data data.yaml \
    --lr0 0.001

Freezing the first 10 layers (the backbone) for initial training preserves the pretrained feature extractors while allowing the detection head to adapt to your dataset. This is especially effective when your dataset is small (<5,000 images).

Export to ONNX

The first step toward edge deployment is ONNX export:

# Export to ONNX
python export.py \
    --weights runs/train/exp/weights/best.pt \
    --include onnx \
    --img 640 \
    --batch 1 \
    --simplify \
    --opset 12

The --simplify flag runs ONNX Simplifier, which fuses operations and optimizes the graph for inference. This typically reduces the graph size by 30-40% and improves inference speed.

Verifying ONNX Export

import onnxruntime as ort
import numpy as np
import cv2

# Load ONNX model
session = ort.InferenceSession("best.onnx", providers=["CPUExecutionProvider"])

# Prepare input
image = cv2.imread("test_image.jpg")
image = cv2.resize(image, (640, 640))
image = image.transpose(2, 0, 1).astype(np.float32) / 255.0
image = np.expand_dims(image, axis=0)

# Run inference
outputs = session.run(None, {"images": image})
print(f"Output shape: {outputs[0].shape}")  # [1, 25200, 10]

TensorRT Export for NVIDIA Edge Devices

For NVIDIA Jetson devices (Nano, Xavier, Orin), TensorRT provides the best inference performance:

# Export to TensorRT
python export.py \
    --weights best.pt \
    --include engine \
    --img 640 \
    --batch 1 \
    --device 0 \
    --half  # FP16 quantization

Manual TensorRT Export with Fine-Grained Control

import tensorrt as trt

logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)
parser = trt.OnnxParser(network, logger)

# Parse ONNX
with open("best.onnx", "rb") as f:
    assert parser.parse(f.read())

# Configure builder
config = builder.create_builder_config()
config.max_workspace_size = 2 << 30  # 2GB
config.set_flag(trt.BuilderFlag.FP16)

# Optimize for specific input size
profile = builder.create_optimization_profile()
profile.set_shape("images", (1, 3, 640, 640), (1, 3, 640, 640), (1, 3, 640, 640))
config.add_optimization_profile(profile)

# Build engine
engine = builder.build_engine(network, config)

# Save engine
with open("best.engine", "wb") as f:
    f.write(engine.serialize())

INT8 Quantization for Maximum Speed

For Jetson Nano and other resource-constrained devices, INT8 quantization can provide an additional 2-3x speedup:

import tensorrt as trt

# Create calibrator for INT8
class CalibrationData(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibration_images, batch_size=8):
        self.batch_size = batch_size
        self.images = self._load_images(calibration_images)
        self.current_index = 0
        self.device_input = None

    def _load_images(self, image_paths):
        images = []
        for path in image_paths:
            img = cv2.imread(path)
            img = cv2.resize(img, (640, 640))
            img = img.transpose(2, 0, 1).astype(np.float32) / 255.0
            images.append(img)
        return np.array(images)

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        if self.current_index >= len(self.images):
            return None
        batch = self.images[self.current_index:self.current_index + self.batch_size]
        self.current_index += self.batch_size
        if self.device_input is None:
            self.device_input = cuda.mem_alloc(batch.nbytes)
        cuda.memcpy_htod(self.device_input, batch.tobytes())
        return [int(self.device_input)]

# Configure INT8
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = CalibrationData(calibration_images)

Edge Deployment Patterns

Jetson Deployment

import cv2
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda

class YOLOv5TRT:
    def __init__(self, engine_path, conf_threshold=0.5, iou_threshold=0.45):
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold

        # Load TensorRT engine
        logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, "rb") as f:
            runtime = trt.Runtime(logger)
            self.engine = runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()

        # Allocate buffers
        self.inputs = []
        self.outputs = []
        self.bindings = []
        for binding in self.engine:
            shape = self.engine.get_tensor_shape(binding)
            size = trt.volume(shape)
            dtype = trt.nptype(self.engine.get_tensor_dtype(binding))
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            self.bindings.append(int(device_mem))
            if self.engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT:
                self.inputs.append({"host": host_mem, "device": device_mem})
            else:
                self.outputs.append({"host": host_mem, "device": device_mem})

    def preprocess(self, image):
        """Preprocess image for inference."""
        img = cv2.resize(image, (640, 640))
        img = img.transpose(2, 0, 1).astype(np.float32) / 255.0
        return np.ascontiguousarray(img)

    def postprocess(self, output, orig_shape):
        """Post-process detections."""
        predictions = output.reshape(-1, 85)  # 80 classes + 5 box params

        # Filter by confidence
        obj_conf = predictions[:, 4]
        mask = obj_conf > self.conf_threshold
        predictions = predictions[mask]

        if len(predictions) == 0:
            return []

        # Get class scores
        class_scores = predictions[:, 5:]
        class_ids = np.argmax(class_scores, axis=1)
        scores = obj_conf[mask] * class_scores[np.arange(len(class_scores)), class_ids]

        # Convert boxes from center to corner format
        boxes = predictions[:, :4]
        boxes[:, 0] -= boxes[:, 2] / 2  # x1
        boxes[:, 1] -= boxes[:, 3] / 2  # y1
        boxes[:, 2] += boxes[:, 0]       # x2
        boxes[:, 3] += boxes[:, 1]       # y2

        # Scale to original image size
        scale_x = orig_shape[1] / 640
        scale_y = orig_shape[0] / 640
        boxes[:, [0, 2]] *= scale_x
        boxes[:, [1, 3]] *= scale_y

        # NMS
        indices = cv2.dnn.NMSBoxes(
            boxes.tolist(), scores.tolist(),
            self.conf_threshold, self.iou_threshold
        )

        results = []
        for i in indices:
            results.append({
                "bbox": boxes[i].tolist(),
                "score": float(scores[i]),
                "class_id": int(class_ids[i]),
            })
        return results

    def detect(self, image):
        """Run detection on a single image."""
        orig_shape = image.shape[:2]
        input_data = self.preprocess(image)

        # Copy input to device
        np.copyto(self.inputs[0]["host"], input_data.ravel())
        cuda.memcpy_htod(self.inputs[0]["device"], self.inputs[0]["host"])

        # Run inference
        self.context.execute_v2(bindings=self.bindings)

        # Copy output from device
        cuda.memcpy_dtoh(self.outputs[0]["host"], self.outputs[0]["device"])

        return self.postprocess(self.outputs[0]["host"], orig_shape)

OpenVINO for Intel Edge Devices

# Export to OpenVINO
python export.py \
    --weights best.pt \
    --include openvino \
    --img 640 \
    --half

from openvino.runtime import Core

core = Core()
model = core.read_model("best_openvino_model/best.xml")
compiled_model = core.compile_model(model, "CPU")

# Inference
input_layer = compiled_model.input(0)
output_layer = compiled_model.output(0)

result = compiled_model([input_data])[output_layer]

Accuracy vs. Speed Tradeoffs

The key to successful edge deployment is choosing the right model size for your hardware:

Model	Parameters	mAP@0.5	mAP@0.5:0.95	FPS (T4)	FPS (Jetson Nano)
YOLOv5n	1.9M	45.7	28.0	350	45
YOLOv5s	7.2M	56.8	37.4	250	25
YOLOv5m	21.2M	64.1	45.2	150	12
YOLOv5l	46.5M	67.8	49.0	90	5
YOLOv5x	86.7M	68.9	50.7	55	2

My decision framework:

1. YOLOv5n for Jetson Nano / Raspberry Pi: When inference speed is paramount and objects are large and distinct.
2. YOLOv5s for Jetson Xavier / Intel NUC: The sweet spot for most edge applications—good accuracy with acceptable speed.
3. YOLOv5m for Jetson Orin / GPU-equipped edge: When accuracy is critical and hardware budget allows.
4. YOLOv5l/x for server-side inference: When running on cloud GPUs or powerful edge servers.

Model Distillation

When you need better accuracy than YOLOv5s but faster inference than YOLOv5m, knowledge distillation can help:

import torch
import torch.nn as nn

class DistillationLoss(nn.Module):
    def __init__(self, temperature=20, alpha=0.5):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha
        self.kl_loss = nn.KLDivLoss(reduction="batchmean")

    def forward(self, student_logits, teacher_logits, targets, hard_loss):
        # Soft target loss
        soft_student = F.log_softmax(student_logits / self.temperature, dim=1)
        soft_teacher = F.softmax(teacher_logits / self.temperature, dim=1)
        soft_loss = self.kl_loss(soft_student, soft_teacher) * (self.temperature ** 2)

        # Combined loss
        return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

Train YOLOv5s as the student with YOLOv5m as the teacher. This typically recovers 50-70% of the mAP gap between the two models while maintaining YOLOv5s inference speed.

Production Pipeline

Camera Integration

import cv2
import threading
import queue

class EdgeDetector:
    def __init__(self, engine_path, camera_id=0):
        self.model = YOLOv5TRT(engine_path)
        self.cap = cv2.VideoCapture(camera_id)
        self.frame_queue = queue.Queue(maxsize=2)
        self.result_queue = queue.Queue(maxsize=2)

    def capture_loop(self):
        """Continuous frame capture."""
        while True:
            ret, frame = self.cap.read()
            if not ret:
                break
            if not self.frame_queue.full():
                self.frame_queue.put(frame)

    def inference_loop(self):
        """Continuous inference."""
        while True:
            frame = self.frame_queue.get()
            detections = self.model.detect(frame)
            if not self.result_queue.full():
                self.result_queue.put((frame, detections))

    def run(self):
        """Run detection pipeline."""
        capture_thread = threading.Thread(target=self.capture_loop, daemon=True)
        inference_thread = threading.Thread(target=self.inference_loop, daemon=True)
        capture_thread.start()
        inference_thread.start()

        while True:
            frame, detections = self.result_queue.get()
            # Process detections (alert, log, display)
            self.process_detections(frame, detections)

    def process_detections(self, frame, detections):
        for det in detections:
            if det["score"] > 0.7:
                x1, y1, x2, y2 = map(int, det["bbox"])
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(
                    frame, f"{det['class_id']}: {det['score']:.2f}",
                    (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2
                )

Edge-to-Cloud Sync

import boto3
import json
import time

class EdgeCloudSync:
    def __init__(self, bucket_name, endpoint_url=None):
        self.s3 = boto3.client("s3", endpoint_url=endpoint_url)
        self.bucket = bucket_name
        self.buffer = []
        self.buffer_size = 100

    def log_detection(self, image_path, detections):
        """Buffer detections and sync to cloud periodically."""
        self.buffer.append({
            "timestamp": time.time(),
            "image": image_path,
            "detections": detections,
        })

        if len(self.buffer) >= self.buffer_size:
            self.sync()

    def sync(self):
        """Upload buffered detections to S3."""
        if not self.buffer:
            return

        key = f"detections/{int(time.time())}.json"
        self.s3.put_object(
            Bucket=self.bucket,
            Key=key,
            Body=json.dumps(self.buffer),
        )
        self.buffer.clear()

Lessons from Edge Deployments

1. Profile on target hardware: Models that run at 200 FPS on a desktop GPU may run at 10 FPS on a Jetson Nano. Always profile on the actual deployment hardware.

2. Warm up the engine: TensorRT engines need a few inference passes to stabilize. Run 10-20 warmup inferences before measuring latency.

3. Use FP16, not FP32: On modern NVIDIA hardware, FP16 inference is 2-3x faster than FP32 with negligible accuracy loss. Always use --half for export.

4. Batch size 1 is the norm: Edge deployments almost always process one image at a time. Optimize for this: use dynamic batch size in TensorRT, avoid batch normalization issues.

5. Monitor thermal throttling: Edge devices throttle under sustained load. Monitor GPU/CPU temperature and implement backoff strategies if needed.

6. Design for intermittent connectivity: Edge devices often have unreliable network connections. Buffer results locally and sync when connectivity is available.

7. Test with production lighting: The single biggest cause of accuracy degradation in edge deployments is lighting conditions that differ from training data. Collect training data under production lighting.

Conclusion

YOLOv5 is not the newest or flashiest object detection architecture, but it is the most production-ready option for edge deployment. Its clean codebase, comprehensive export pipeline, and scalable model sizes make it the practical choice when you need to ship a detection model to a device with real constraints.

The key to success is not the model itself—it is the engineering around it. Proper dataset preparation, hardware-aware model selection, optimized export pipelines, and robust deployment infrastructure. Get these right, and YOLOv5 will serve you well from prototype to production.