YOLO Advanced Optimization: Lightweight, Quantization and Accuracy

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

# Large model as teacher, small model as student
teacher = YOLO("yolo26x.pt")
student = YOLO("yolo26n.yaml")

# Distillation training (Ultralytics built-in support)
student.train(
    data="data.yaml",
    distill="yolo26x.pt",  # Teacher model
    distill_ratio=0.5,     # Distillation loss ratio
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

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

# L1 unstructured pruning on conv layers
model = YOLO("yolo26n.pt")
for name, module in model.model.named_modules():
    if isinstance(module, torch.nn.Conv2d):
        prune.l1_unstructured(module, name="weight", amount=0.3)
        prune.remove(module, "weight")  # Make pruning permanent

# Channel pruning with torch-pruning library
# pip install torch-pruning
import torch_pruning as tp

model = YOLO("yolo26n.pt").model
DG = tp.DependencyGraph()
DG.build_dependency(model, example_inputs=torch.randn(1, 3, 640, 640))

# Prune 20% channels by L1 norm
pruning_plan = DG.get_pruning_plan(
    model.model[4], tp.prune_conv,
    pruning_dim=0,  # Output channel dimension
    idxs=list(range(0, 64, 5))  # Keep every 5th channel
)
pruning_plan.exec()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

model = YOLO("yolo26n.pt")

# INT8 quantization (requires calibration data)
model.export(
    format="engine",      # TensorRT
    int8=True,
    data="data.yaml",     # Calibration dataset
    batch=8,
)

# ONNX dynamic quantization
model.export(
    format="onnx",
    dynamic=True,
    simplify=True,
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

import tensorrt as trt
from ctypes import c_size_t

class YOLOCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, data_loader, cache_file="calibration.cache"):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.data_loader = data_loader
        self.cache_file = cache_file
        self.buffer_size = 0

    def get_batch_size(self):
        return 8

    def get_batch(self, names):
        try:
            batch = next(self.data_loader)
            batch = batch.cpu().numpy()
            return [batch.astype(np.float32).ctypes.data_as(c_size_t)]
        except StopIteration:
            return None

1
2

Input → QuantizeLinear → Conv → DequantizeLinear → ... → Output
          scale, zp                scale, zp

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

from onnxruntime.quantization import quantize_dynamic, quantize_static
from onnxruntime.quantization import CalibrationMethod, QuantType, QuantFormat

# Dynamic quantization (no calibration data needed)
quantize_dynamic(
    "yolo26n.onnx",
    "yolo26n_dynamic.onnx",
    weight_type=QuantType.QInt8,
)

# Static quantization (requires calibration data)
class YOLOCalibDataReader:
    def __init__(self, image_paths, batch_size=8):
        self.images = [self.preprocess(img) for img in image_paths]
        self.batch_size = batch_size
        self.index = 0

    def get_next(self):
        if self.index >= len(self.images):
            return None
        batch = self.images[self.index:self.index + self.batch_size]
        self.index += self.batch_size
        return {"images": np.stack(batch)}

quantize_static(
    "yolo26n.onnx",
    "yolo26n_static.onnx",
    calibration_data_reader=YOLOCalibDataReader(calibration_images),
    quant_format=QuantFormat.QDQ,  # QDQ format, supports more optimizations
    per_channel=True,              # Per-channel gives higher accuracy
    activation_type=QuantType.QInt8,
    weight_type=QuantType.QInt8,
    calibrate_method=CalibrationMethod.Entropy,
)

1
2
3
4
5
6

# Convert ONNX to NCNN format
onnx2ncnn yolo26n.onnx yolo26n.param yolo26n.bin

# Use Vulkan GPU for accelerated inference
ncnnoptimize yolo26n.param yolo26n.bin yolo26n_opt.param yolo26n_opt.bin 1
# Last parameter: 0=CPU, 1=Vulkan

1
2
3
4
5
6
7

# FP16 storage mode
ncnn2table --param=yolo26n.param --bin=yolo26n.bin \
    --input=calibration_images --output=table.bfp32 \
    --mean=0,0,0 --norm=255,255,255 --size=640,640

ncnn2int8 yolo26n.param yolo26n.bin yolo26n_int8.param \
    yolo26n_int8.bin table.bfp32

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

// C++ integration example
#include "net.h"

ncnn::Net yolo;
yolo.load_param("yolo26n_opt.param");
yolo.load_model("yolo26n_opt.bin");

ncnn::Mat in = ncnn::Mat::from_pixels_resize(
    image_data, ncnn::Mat::PIXEL_BGR, 
    w, h, 640, 640
);

ncnn::Extractor ex = yolo.create_extractor();
ex.input("images", in);
ncnn::Mat out;
ex.extract("output", out);

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

# Multi-scale during training
model.train(
    imgsz=640,
    multi_scale=True,  # Auto vary ±50% range
)

# Multi-scale enhancement during inference
results = model(
    "test.jpg",
    imgsz=[640, 800, 1024],  # Multi-scale fusion
    augment=True,            # TTA augmentation during testing
)

1
2
3
4
5

model.train(
    box=7.5,    # Box regression weight
    cls=0.5,    # Classification weight (reduce for minority classes)
    dfl=1.5,
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

import torch.nn.functional as F

class FocalLoss(torch.nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, pred, target):
        ce_loss = F.binary_cross_entropy_with_logits(
            pred, target, reduction="none"
        )
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
        return focal_loss.mean()

# Replace loss function in training callback
# Ultralytics custom loss: inherit DetectionLoss and override bcecls

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

import numpy as np

def compute_class_weights(labels_path, num_classes=80):
    """Compute class weights based on annotation frequency"""
    class_counts = np.zeros(num_classes)
    # Count bounding boxes per class
    for label_file in labels_path.glob("*.txt"):
        boxes = np.loadtxt(label_file)  # [class_id, x, y, w, h]
        if boxes.ndim == 1:
            boxes = boxes.reshape(1, -1)
        classes = boxes[:, 0].astype(int)
        for c in classes:
            class_counts[c] += 1

    # Median Frequency Balancing
    median_freq = np.median(class_counts[class_counts > 0])
    class_weights = median_freq / (class_counts + 1e-6)
    class_weights = np.clip(class_weights, 0.1, 10.0)  # Clamp range
    return class_weights

# weights = compute_class_weights("datasets/coco/labels/train/")
# Pass weights into loss function

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

# Resampling: set sampling weights in data.yaml
# Or use torch's WeightedRandomSampler

from torch.utils.data import WeightedRandomSampler
from collections import Counter

def create_balanced_sampler(labels_path, num_classes=80):
    class_counts = Counter()
    for label_file in labels_path.glob("*.txt"):
        boxes = np.loadtxt(label_file)
        if boxes.ndim == 1:
            boxes = boxes.reshape(1, -1)
        for c in boxes[:, 0].astype(int):
            class_counts[c] += 1

    # Sample weight = 1 / class frequency
    weights = [1.0 / class_counts[c] for c in range(num_classes)]

    # Per-image weight = average weight of all its annotation classes
    sample_weights = []
    for label_file in labels_path.glob("*.txt"):
        boxes = np.loadtxt(label_file)
        if boxes.ndim == 0:
            sample_weights.append(1.0)
        elif boxes.ndim == 1:
            sample_weights.append(weights[int(boxes[0])])
        else:
            sample_weights.append(np.mean(
                [weights[int(c)] for c in boxes[:, 0]]
            ))

    return WeightedRandomSampler(
        sample_weights, len(sample_weights), replacement=True
    )

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

import time
import numpy as np
import torch

def benchmark_latency(model, input_tensor, num_warmup=50, num_runs=200):
    # Warmup
    for _ in range(num_warmup):
        _ = model(input_tensor)

    # Formal measurement
    if torch.cuda.is_available():
        torch.cuda.synchronize()

    timings = []
    for _ in range(num_runs):
        start = time.perf_counter()
        _ = model(input_tensor)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        timings.append((time.perf_counter() - start) * 1000)  # ms

    return {
        "mean": np.mean(timings),
        "std": np.std(timings),
        "p50": np.percentile(timings, 50),
        "p95": np.percentile(timings, 95),
        "p99": np.percentile(timings, 99),
    }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

def benchmark_throughput(
    model, batch_size=32, input_size=(3, 640, 640), num_batches=100
):
    dummy_input = torch.randn(batch_size, *input_size).cuda()

    # Warmup
    for _ in range(10):
        _ = model(dummy_input)

    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(num_batches):
        _ = model(dummy_input)
    torch.cuda.synchronize()
    total_time = time.perf_counter() - start

    total_images = batch_size * num_batches
    throughput = total_images / total_time
    return {
        "throughput": throughput,
        "total_time": total_time,
        "batch_size": batch_size,
    }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

from torch.profiler import profile, record_function, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    with record_function("model_inference"):
        output = model(input_tensor)

# Print time ranking
print(prof.key_averages().table(
    sort_by="cuda_time_total", row_limit=20
))

# Export for Chrome Trace visualization
prof.export_chrome_trace("trace.json")
# View in chrome://tracing/

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

import onnxruntime as ort

# Enable session profiling
sess_options = ort.SessionOptions()
sess_options.enable_profiling = True
sess = ort.InferenceSession("yolo26n.onnx", sess_options)

# Run inference
outputs = sess.run(None, {"images": np_input})

# Get profiling data
prof_file = sess.end_profiling()
# Outputs profiling_*.json file, viewable in chrome://tracing/

1
2
3
4
5
6

# Install: https://developer.nvidia.com/nsight-systems
# Command-line profiling
nsys profile -o yolo_profile -t cuda,nvtx python inference.py

# Visualize results
# Open yolo_profile.qdrep in Nsight Systems GUI