YOLO性能优化与部署实践：从模型压缩到多平台落地

# 使用PyTorch进行动态量化
import torch
import torch.nn as nn

# 加载预训练模型
model = torch.load('yolov10n.pt')
model.eval()

# 动态量化
quantized_model = torch.quantization.quantize_dynamic(
    model,
    {nn.Linear, nn.Conv2d},
    dtype=torch.qint8
)

# 保存量化模型
torch.save(quantized_model, 'yolov10n_quantized.pt')

# 使用PyTorch进行静态量化
import torch
import torch.nn as nn
import torch.quantization

# 定义量化配置
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# 准备模型
model_prepared = torch.quantization.prepare(model)

# 校准模型（使用校准集）
for data, _ in calibration_dataloader:
    model_prepared(data)

# 转换为量化模型
quantized_model = torch.quantization.convert(model_prepared)

# 保存量化模型
torch.save(quantized_model, 'yolov10n_quantized_static.pt')

# 使用torch.nn.utils.prune进行通道剪枝
import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

# 加载模型
model = torch.load('yolov10n.pt')

# 对卷积层进行通道剪枝
for name, module in model.named_modules():
    if isinstance(module, nn.Conv2d):
        # 移除30%的通道
        prune.l1_unstructured(module, name='weight', amount=0.3)
        # 永久化剪枝
        prune.remove(module, 'weight')

# 保存剪枝后的模型
torch.save(model, 'yolov10n_pruned.pt')

# 知识蒸馏实现示例
import torch
import torch.nn as nn
import torch.optim as optim

# 定义教师模型和学生模型
teacher_model = torch.load('yolov10l.pt')
teacher_model.eval()

student_model = torch.load('yolov10n.pt')
student_model.train()

# 定义损失函数
criterion_cls = nn.CrossEntropyLoss()
criterion_kl = nn.KLDivLoss(reduction='batchmean')

# 定义优化器
optimizer = optim.Adam(student_model.parameters(), lr=1e-4)

# 蒸馏温度
temperature = 5.0

# 蒸馏权重
alpha = 0.7

# 训练循环
def train_distillation(dataloader, teacher_model, student_model, criterion_cls, criterion_kl, optimizer, temperature, alpha):
    student_model.train()
    teacher_model.eval()
    
    for batch, (X, y) in enumerate(dataloader):
        # 教师模型推理（不更新参数）
        with torch.no_grad():
            teacher_pred = teacher_model(X)
        
        # 学生模型推理
        student_pred = student_model(X)
        
        # 计算硬标签损失
        loss_hard = criterion_cls(student_pred, y)
        
        # 计算软标签损失（KL散度）
        loss_soft = criterion_kl(
            torch.nn.functional.log_softmax(student_pred / temperature, dim=1),
            torch.nn.functional.softmax(teacher_pred / temperature, dim=1)
        ) * (temperature ** 2)
        
        # 总损失
        loss = alpha * loss_soft + (1 - alpha) * loss_hard
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{len(dataloader.dataset):>5d}]")

import torch
from ultralytics import YOLO

# 加载模型
model = YOLO('yolov10n.pt')

# 导出为ONNX格式
model.export(format='onnx', imgsz=640, batch=1)

import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

# 加载ONNX模型
onnx_model = onnx.load('yolov10n.onnx')

# 使用ONNX Runtime进行动态量化
quantized_model = quantize_dynamic(
    'yolov10n.onnx',
    'yolov10n_quantized.onnx',
    weight_type=QuantType.QUInt8
)

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

# 创建TensorRT引擎
class YOLOv10TRT:
    def __init__(self, engine_path, imgsz=640):
        self.imgsz = imgsz
        self.logger = trt.Logger(trt.Logger.WARNING)
        
        # 加载引擎
        with open(engine_path, 'rb') as f:
            engine_data = f.read()
        
        runtime = trt.Runtime(self.logger)
        self.engine = runtime.deserialize_cuda_engine(engine_data)
        self.context = self.engine.create_execution_context()
        
        # 分配内存
        self.inputs = []
        self.outputs = []
        self.allocations = []
        
        for i in range(self.engine.num_bindings):
            name = self.engine.get_binding_name(i)
            dtype = trt.nptype(self.engine.get_binding_dtype(i))
            shape = self.engine.get_binding_shape(i)
            
            if self.engine.binding_is_input(i):
                shape = (1, 3, imgsz, imgsz)  # 动态batch
            
            size = np.prod(shape) * dtype().itemsize
            allocation = cuda.mem_alloc(size)
            binding = {}
            binding['name'] = name
            binding['dtype'] = dtype
            binding['shape'] = shape
            binding['allocation'] = allocation
            
            if self.engine.binding_is_input(i):
                self.inputs.append(binding)
            else:
                self.outputs.append(binding)
    
    def infer(self, img):
        # 预处理图像
        img = cv2.resize(img, (self.imgsz, self.imgsz))
        img = img.transpose((2, 0, 1))  # HWC to CHW
        img = img[np.newaxis, ...]  # 添加batch维度
        img = img.astype(np.float32) / 255.0  # 归一化
        
        # 复制数据到设备
        cuda.memcpy_htod(self.inputs[0]['allocation'], img.ravel())
        
        # 执行推理
        self.context.execute_v2([b['allocation'] for b in self.inputs + self.outputs])
        
        # 从设备复制结果
        outputs = []
        for output in self.outputs:
            output_buffer = np.empty(output['shape'], dtype=output['dtype'])
            cuda.memcpy_dtoh(output_buffer, output['allocation'])
            outputs.append(output_buffer)
        
        return outputs

import onnxruntime as ort
import numpy as np
import cv2

# 加载ONNX模型
session = ort.InferenceSession('yolov10n.onnx', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

# 获取输入输出名称
input_name = session.get_inputs()[0].name
output_names = [output.name for output in session.get_outputs()]

# 图像预处理
def preprocess(img, imgsz=640):
    img = cv2.resize(img, (imgsz, imgsz))
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, HWC to CHW
    img = np.ascontiguousarray(img)
    img = img.astype(np.float32) / 255.0
    img = np.expand_dims(img, axis=0)
    return img

# 模型推理
def infer(img, session, input_name, output_names):
    # 预处理
    input_tensor = preprocess(img)
    
    # 推理
    outputs = session.run(output_names, {input_name: input_tensor})
    
    return outputs

# 后处理
def postprocess(outputs, img_shape, conf_thres=0.5, iou_thres=0.45):
    # YOLOv10的输出处理逻辑
    # 这里简化处理，实际需要根据具体输出格式进行调整
    boxes = outputs[0][0]
    scores = outputs[1][0]
    classes = outputs[2][0]
    
    # 过滤低置信度框
    mask = scores > conf_thres
    boxes = boxes[mask]
    scores = scores[mask]
    classes = classes[mask]
    
    # 非极大值抑制
    indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), conf_thres, iou_thres)
    
    # 转换为原始图像坐标
    h, w = img_shape[:2]
    scale = min(640/w, 640/h)
    pad_w = (640 - w * scale) / 2
    pad_h = (640 - h * scale) / 2
    
    result = []
    for i in indices:
        box = boxes[i]
        # 转换为原始图像坐标
        x1 = (box[0] - pad_w) / scale
        y1 = (box[1] - pad_h) / scale
        x2 = (box[2] - pad_w) / scale
        y2 = (box[3] - pad_h) / scale
        
        result.append({
            'box': [x1, y1, x2, y2],
            'score': scores[i],
            'class': classes[i]
        })
    
    return result

from openvino.runtime import Core
import numpy as np
import cv2

# 加载模型
ie = Core()
model = ie.read_model(model='yolov10n.xml')
compiled_model = ie.compile_model(model=model, device_name='CPU')

# 获取输入输出
input_layer = compiled_model.input(0)
output_layer = compiled_model.output(0)

# 图像预处理
def preprocess(img, imgsz=640):
    img = cv2.resize(img, (imgsz, imgsz))
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, HWC to CHW
    img = np.ascontiguousarray(img)
    img = img.astype(np.float32) / 255.0
    img = np.expand_dims(img, axis=0)
    return img

# 模型推理
def infer(img, compiled_model, input_layer, output_layer):
    # 预处理
    input_tensor = preprocess(img)
    
    # 推理
    outputs = compiled_model([input_tensor])[output_layer]
    
    return outputs

import torch
import tensorflow as tf
from ultralytics import YOLO

# 1. 导出为ONNX
model = YOLO('yolov10n.pt')
model.export(format='onnx', imgsz=640, batch=1)

# 2. 转换为TFLite
import onnx
from onnx_tf.backend import prepare

# 加载ONNX模型
onnx_model = onnx.load('yolov10n.onnx')

# 转换为TensorFlow模型
tf_rep = prepare(onnx_model)
tf_rep.export_graph('yolov10n.pb')

# 转换为TFLite模型
converter = tf.lite.TFLiteConverter.from_saved_model('yolov10n.pb')
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

# 保存TFLite模型
with open('yolov10n.tflite', 'wb') as f:
    f.write(tflite_model)

# 3. TFLite推理
interpreter = tf.lite.Interpreter(model_path='yolov10n.tflite')
interpreter.allocate_tensors()

# 获取输入输出张量
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 推理函数
def tflite_infer(img, interpreter, input_details, output_details):
    # 预处理
    input_tensor = preprocess(img)
    
    # 设置输入
    interpreter.set_tensor(input_details[0]['index'], input_tensor)
    
    # 推理
    interpreter.invoke()
    
    # 获取输出
    outputs = []
    for output_detail in output_details:
        output = interpreter.get_tensor(output_detail['index'])
        outputs.append(output)
    
    return outputs

from ultralytics import YOLO

# 加载模型
model = YOLO('yolov10n.pt')

# 自动优化与部署
# 导出为ONNX
model.export(format='onnx', imgsz=640, batch=1, optimize=True)

# 导出为TensorRT
model.export(format='engine', imgsz=640, batch=1, device=0)

# 导出为OpenVINO
model.export(format='openvino', imgsz=640, batch=1)

# 导出为TFLite
model.export(format='tflite', imgsz=640, batch=1)

import time
import numpy as np
import cv2
from ultralytics import YOLO

# 加载不同版本的模型
models = {
    'yolov10n': YOLO('yolov10n.pt'),
    'yolov10n_quantized': YOLO('yolov10n_quantized.pt'),
    'yolov10n_pruned': YOLO('yolov10n_pruned.pt')
}

# 加载测试图像
img = cv2.imread('test.jpg')

# 性能测试
def benchmark(model, img, iterations=100):
    # 预热
    for _ in range(10):
        model(img)
    
    # 测试
    start_time = time.time()
    for _ in range(iterations):
        results = model(img)
    end_time = time.time()
    
    # 计算平均推理时间
    avg_time = (end_time - start_time) / iterations * 1000  # 毫秒
    fps = 1000 / avg_time
    
    return avg_time, fps

# 运行测试
for name, model in models.items():
    avg_time, fps = benchmark(model, img)
    print(f"{name}: 平均推理时间 {avg_time:.2f} ms, FPS {fps:.2f}")

# 1. 模型量化
python -m ultralytics export --model yolov10n.pt --format onnx --quantize dynamic

# 2. 模型剪枝
python -m ultralytics prune --model yolov10n.pt --prune-ratio 0.3

# 3. 知识蒸馏
python -m ultralytics distill --teacher-model yolov10l.pt --student-model yolov10n.pt --data data.yaml --epochs 50

# 4. 导出到不同格式
python -m ultralytics export --model yolov10n.pt --format onnx
python -m ultralytics export --model yolov10n.pt --format engine --device 0
python -m ultralytics export --model yolov10n.pt --format openvino
python -m ultralytics export --model yolov10n.pt --format tflite

# TensorRT部署环境
pip install tensorrt
pip install pycuda

# ONNX Runtime部署环境
pip install onnxruntime-gpu  # GPU版本
pip install onnxruntime  # CPU版本

# OpenVINO部署环境
pip install openvino

# TFLite部署环境
pip install tensorflow
pip install tflite-runtime

# 安装YOLO
pip install ultralytics

import time
import numpy as np
import cv2
from ultralytics import YOLO

# 加载测试图像
img = cv2.imread('test.jpg')
if img is None:
    print("无法加载测试图像")
    exit()

# 模型列表
model_paths = [
    'yolov10n.pt',
    'yolov10s.pt', 
    'yolov10m.pt',
    'yolov10l.pt',
    'yolov10x.pt'
]

# 性能测试函数
def benchmark_model(model_path, img, iterations=100):
    # 加载模型
    model = YOLO(model_path)
    
    # 预热
    for _ in range(10):
        results = model(img)
    
    # 测试
    start_time = time.time()
    for _ in range(iterations):
        results = model(img)
    end_time = time.time()
    
    # 计算性能指标
    avg_time = (end_time - start_time) / iterations * 1000  # 毫秒
    fps = 1000 / avg_time
    
    # 获取模型信息
    model_name = model_path.split('.')[0]
    
    return {
        'model': model_name,
        'avg_time': avg_time,
        'fps': fps
    }

# 运行测试
results = []
for model_path in model_paths:
    result = benchmark_model(model_path, img)
    results.append(result)
    print(f"{result['model']}: 平均推理时间 {result['avg_time']:.2f} ms, FPS {result['fps']:.2f}")

# 保存结果到CSV
import csv
with open('yolo_performance.csv', 'w', newline='') as f:
    writer = csv.DictWriter(f, fieldnames=['model', 'avg_time', 'fps'])
    writer.writeheader()
    for result in results:
        writer.writerow(result)