YOLO 进阶优化：轻量化、量化与精度提升

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

# 大模型作为教师，小模型作为学生
teacher = YOLO("yolo26x.pt")
student = YOLO("yolo26n.yaml")

# 蒸馏训练（Ultralytics内置支持）
student.train(
    data="data.yaml",
    distill="yolo26x.pt",  # 教师模型
    distill_ratio=0.5,     # 蒸馏损失比例
)

 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 非结构化剪枝
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")  # 使剪枝永久化

# 通道剪枝使用 torch-pruning 库
# 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))

# 按 L1 范数剪枝 20% 通道
pruning_plan = DG.get_pruning_plan(
    model.model[4], tp.prune_conv,
    pruning_dim=0,  # 输出通道维度
    idxs=list(range(0, 64, 5))  # 每 5 个通道保留一个
)
pruning_plan.exec()

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

model = YOLO("yolo26n.pt")

# INT8量化（需要校准数据）
model.export(
    format="engine",      # TensorRT
    int8=True,
    data="data.yaml",     # 校准数据集
    batch=8,
)

# ONNX动态量化
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

输入 → QuantizeLinear → Conv → DequantizeLinear → ... → 输出
          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

# 动态量化（无需校准数据）
quantize_dynamic(
    "yolo26n.onnx",
    "yolo26n_dynamic.onnx",
    weight_type=QuantType.QInt8,
)

# 静态量化（需要校准数据）
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 格式，支持更多优化
    per_channel=True,              # 逐通道量化精度更高
    activation_type=QuantType.QInt8,
    weight_type=QuantType.QInt8,
    calibrate_method=CalibrationMethod.Entropy,
)

1
2
3
4
5
6

# 将 ONNX 转换为 NCNN 格式
onnx2ncnn yolo26n.onnx yolo26n.param yolo26n.bin

# 使用 Vulkan GPU 加速推理
ncnnoptimize yolo26n.param yolo26n.bin yolo26n_opt.param yolo26n_opt.bin 1
# 最后一个参数: 0=CPU, 1=Vulkan

1
2
3
4
5
6
7

# FP16 存储模式
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++ 集成示例
#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

# 训练时多尺度
model.train(
    imgsz=640,
    multi_scale=True,  # 自动在 ±50% 范围内变化
)

# 推理时多尺度增强
results = model(
    "test.jpg",
    imgsz=[640, 800, 1024],  # 多尺度融合
    augment=True,            # TTA测试时增强
)

1
2
3
4
5

model.train(
    box=7.5,    # 框回归权重
    cls=0.5,    # 分类权重（少类别时降低）
    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()

# 在训练回调中替换损失函数
# Ultralytics 自定义损失：继承 DetectionLoss 并重写 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):
    """根据标注频率计算类别权重"""
    class_counts = np.zeros(num_classes)
    # 统计每个类别的标注框数量
    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)  # 限制范围
    return class_weights

# weights = compute_class_weights("datasets/coco/labels/train/")
# 将 weights 传入损失函数

 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

# 重采样：在 data.yaml 中设置采样权重
# 或者使用 torch 的 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

    # 样本权重 = 1 / 类别频率
    weights = [1.0 / class_counts[c] for c in range(num_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)

    # 正式测量
    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(prof.key_averages().table(
    sort_by="cuda_time_total", row_limit=20
))

# 导出为 Chrome Trace 可视化
prof.export_chrome_trace("trace.json")
# 在 chrome://tracing/ 中打开

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

import onnxruntime as ort

# 启用会话 profiling
sess_options = ort.SessionOptions()
sess_options.enable_profiling = True
sess = ort.InferenceSession("yolo26n.onnx", sess_options)

# 运行推理
outputs = sess.run(None, {"images": np_input})

# 获取 profile 数据
prof_file = sess.end_profiling()
# 输出 profiling_*.json 文件，可在 chrome://tracing/ 中查看

1
2
3
4
5
6

# 安装: https://developer.nvidia.com/nsight-systems
# 命令行 profiling
nsys profile -o yolo_profile -t cuda,nvtx python inference.py

# 可视化结果
# 在 Nsight Systems GUI 中打开 yolo_profile.qdrep