第14章 视觉模型部署实战¶

import torch
import torchvision

# 加载模型
# torchvision 0.13+ 推荐使用枚举类型而非字符串
from torchvision.models import ResNet50_Weights
model = torchvision.models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
# 旧版本兼容写法：model = torchvision.models.resnet50(weights="IMAGENET1K_V2")
model.eval()  # eval()评估模式

# 创建虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 基础导出
torch.onnx.export(
    model,
    dummy_input,
    "resnet50.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=17,                    # 使用较新的opset
    dynamic_axes={                       # 动态batch
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    },
    do_constant_folding=True,            # 常量折叠优化
)
print("✅ ONNX导出成功")

from ultralytics import YOLO

# 加载YOLOv8模型
model = YOLO("yolov8n.pt")  # n/s/m/l/x

# 导出ONNX（带NMS后处理）
model.export(
    format="onnx",
    imgsz=640,
    dynamic=True,        # 动态batch
    simplify=True,       # onnx-simplifier优化
    opset=17,
    half=False,          # FP16权重
)

import onnx
import onnxruntime as ort
import numpy as np

# 1. 模型验证
model = onnx.load("resnet50.onnx")
onnx.checker.check_model(model)
print(f"✅ ONNX模型验证通过")
print(f"  输入: {[i.name for i in model.graph.input]}")
print(f"  输出: {[o.name for o in model.graph.output]}")

# 2. 模型简化（消除冗余节点）
import onnxsim
model_simplified, check = onnxsim.simplify(model)
assert check, "简化失败"  # assert断言
onnx.save(model_simplified, "resnet50_sim.onnx")

# 3. 精度验证：对比PyTorch和ONNX输出
session = ort.InferenceSession("resnet50.onnx")
dummy_np = np.random.randn(1, 3, 224, 224).astype(np.float32)

# ONNX Runtime推理
ort_output = session.run(None, {"input": dummy_np})[0]

# PyTorch推理
import torchvision
pt_model = torchvision.models.resnet50(weights="IMAGENET1K_V2")
pt_model.eval()
with torch.no_grad():  # 禁用梯度计算，节省内存
    pt_output = pt_model(torch.from_numpy(dummy_np)).numpy()

# 对比
max_diff = np.abs(ort_output - pt_output).max()
print(f"  最大精度差异: {max_diff:.6f}")
assert max_diff < 1e-5, "精度验证失败！"
print("✅ 精度验证通过")

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

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_path, engine_path, fp16=True, int8=False,
                 max_batch=16, workspace_gb=2):
    """从ONNX构建TensorRT Engine"""
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, TRT_LOGGER)

    # 解析ONNX
    with open(onnx_path, 'rb') as f:  # with自动管理文件关闭
        if not parser.parse(f.read()):
            for i in range(parser.num_errors):
                print(f"Error: {parser.get_error(i)}")
            return None

    # 构建配置
    config = builder.create_builder_config()
    config.set_memory_pool_limit(
        trt.MemoryPoolType.WORKSPACE, workspace_gb << 30)

    if fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    if int8:
        config.set_flag(trt.BuilderFlag.INT8)
        # 需要设置校准器（参见INT8校准部分）

    # 动态shape profile
    profile = builder.create_optimization_profile()
    profile.set_shape("input",
                      min=(1, 3, 224, 224),
                      opt=(8, 3, 224, 224),
                      max=(max_batch, 3, 224, 224))
    config.add_optimization_profile(profile)

    # 构建Engine
    print("Building TensorRT engine (this may take a few minutes)...")
    serialized = builder.build_serialized_network(network, config)

    # 保存
    with open(engine_path, 'wb') as f:
        f.write(serialized)

    print(f"✅ Engine saved to {engine_path}")
    return serialized

class TRTInference:
    """
    TensorRT推理封装

    **版本要求**：
    - TensorRT >= 8.6.0（使用 num_io_tensors API）
    - 对于 TensorRT < 8.6，请使用 TRTInferenceLegacy 类

    **API变更说明**：
    - TensorRT 8.6+ 引入了 `num_io_tensors` 和 `get_tensor_name` API
    - 早期版本使用 `num_bindings` 和 `get_binding_name` API
    """

    def __init__(self, engine_path):
        self.logger = trt.Logger(trt.Logger.WARNING)

        # 反序列化Engine
        runtime = trt.Runtime(self.logger)
        with open(engine_path, 'rb') as f:
            self.engine = runtime.deserialize_cuda_engine(f.read())

        self.context = self.engine.create_execution_context()
        self.stream = cuda.Stream()

        # 分析输入输出（TensorRT 8.6+ API）
        self.inputs = []
        self.outputs = []
        self.bindings = {}

        for i in range(self.engine.num_io_tensors):
            name = self.engine.get_tensor_name(i)
            dtype = trt.nptype(self.engine.get_tensor_dtype(name))
            mode = self.engine.get_tensor_mode(name)

            if mode == trt.TensorIOMode.INPUT:
                self.inputs.append({"name": name, "dtype": dtype})
            else:
                self.outputs.append({"name": name, "dtype": dtype})


class TRTInferenceLegacy:
    """
    TensorRT推理封装（兼容 TensorRT 8.5 及更早版本）

    **适用版本**：TensorRT < 8.6.0
    **主要差异**：使用 `num_bindings` 和 `get_binding_name` API
    """

    def __init__(self, engine_path):
        self.logger = trt.Logger(trt.Logger.WARNING)

        # 反序列化Engine
        runtime = trt.Runtime(self.logger)
        with open(engine_path, 'rb') as f:
            self.engine = runtime.deserialize_cuda_engine(f.read())

        self.context = self.engine.create_execution_context()
        self.stream = cuda.Stream()

        # 分析输入输出（TensorRT 8.5 及更早版本 API）
        self.inputs = []
        self.outputs = []
        self.bindings = []

        for i in range(self.engine.num_bindings):
            name = self.engine.get_binding_name(i)
            dtype = trt.nptype(self.engine.get_binding_dtype(i))

            if self.engine.binding_is_input(i):
                self.inputs.append({"name": name, "dtype": dtype, "index": i})
            else:
                self.outputs.append({"name": name, "dtype": dtype, "index": i})

    def infer(self, input_data: np.ndarray) -> np.ndarray:
        """执行推理（兼容旧版API）"""
        # 分配GPU内存
        d_input = cuda.mem_alloc(input_data.nbytes)

        # 获取输出shape
        output_info = self.outputs[0]
        output_shape = self.context.get_binding_shape(output_info["index"])
        output_data = np.empty(output_shape, dtype=output_info["dtype"])
        d_output = cuda.mem_alloc(output_data.nbytes)

        # 设置bindings
        self.bindings = [None] * self.engine.num_bindings
        self.bindings[self.inputs[0]["index"]] = int(d_input)
        self.bindings[output_info["index"]] = int(d_output)

        # 拷贝输入到GPU
        cuda.memcpy_htod_async(d_input, input_data.ravel(), self.stream)

        # 执行推理
        self.context.execute_async_v2(self.bindings, self.stream.handle)

        # 拷贝输出到CPU
        cuda.memcpy_dtoh_async(output_data, d_output, self.stream)
        self.stream.synchronize()

        # 释放GPU内存
        d_input.free()
        d_output.free()

        return output_data


def create_trt_inference(engine_path):
    """
    自动选择合适的TensorRT推理类

    根据TensorRT版本自动选择 TRTInference 或 TRTInferenceLegacy

    Args:
        engine_path: TensorRT引擎文件路径

    Returns:
        TRTInference 或 TRTInferenceLegacy 实例
    """
    trt_version = tuple(map(int, trt.__version__.split('.')))

    if trt_version >= (8, 6, 0):
        print(f"✅ 使用 TensorRT {trt.__version__} (>= 8.6) API")
        return TRTInference(engine_path)
    else:
        print(f"⚠️ 使用 TensorRT {trt.__version__} (< 8.6) 兼容API")
        return TRTInferenceLegacy(engine_path)

    def infer(self, input_data: np.ndarray) -> np.ndarray:
        """执行推理"""
        batch_size = input_data.shape[0]

        # 设置动态shape
        input_name = self.inputs[0]["name"]
        self.context.set_input_shape(input_name, input_data.shape)

        # 分配GPU内存
        d_input = cuda.mem_alloc(input_data.nbytes)

        # 获取输出shape
        output_name = self.outputs[0]["name"]
        output_shape = self.context.get_tensor_shape(output_name)
        output_data = np.empty(output_shape,
                               dtype=self.outputs[0]["dtype"])
        d_output = cuda.mem_alloc(output_data.nbytes)

        # 拷贝输入到GPU
        cuda.memcpy_htod_async(d_input, input_data.ravel(), self.stream)

        # 设置tensor地址
        self.context.set_tensor_address(input_name, int(d_input))
        self.context.set_tensor_address(output_name, int(d_output))

        # 执行推理
        self.context.execute_async_v3(self.stream.handle)

        # 拷贝输出到CPU
        cuda.memcpy_dtoh_async(output_data, d_output, self.stream)
        self.stream.synchronize()

        # 释放GPU内存
        d_input.free()
        d_output.free()

        return output_data

# 使用示例
if __name__ == "__main__":
    # 构建Engine
    build_engine("resnet50.onnx", "resnet50_fp16.engine", fp16=True)

    # 推理
    infer = TRTInference("resnet50_fp16.engine")

    # 模拟输入
    input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
    output = infer.infer(input_data)

    top5 = np.argsort(output[0])[-5:][::-1]
    print(f"Top-5 classes: {top5}")

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

class YOLOv8TRT:
    """YOLOv8 TensorRT推理"""

    def __init__(self, engine_path, conf_thres=0.5, iou_thres=0.45):
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres

        # 加载Engine
        logger = trt.Logger(trt.Logger.WARNING)
        runtime = trt.Runtime(logger)
        with open(engine_path, 'rb') as f:
            self.engine = runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        self.stream = cuda.Stream()

    def preprocess(self, img, input_size=640):
        """LetterBox预处理"""
        h, w = img.shape[:2]  # 切片操作，取前n个元素
        scale = min(input_size / h, input_size / w)
        new_h, new_w = int(h * scale), int(w * scale)

        # resize
        resized = cv2.resize(img, (new_w, new_h))

        # padding
        canvas = np.full((input_size, input_size, 3), 114, dtype=np.uint8)
        dh, dw = (input_size - new_h) // 2, (input_size - new_w) // 2
        canvas[dh:dh + new_h, dw:dw + new_w] = resized

        # 归一化 + HWC→CHW + 增加batch维
        blob = canvas.astype(np.float32) / 255.0
        blob = blob.transpose(2, 0, 1)[np.newaxis, ...]
        blob = np.ascontiguousarray(blob)

        return blob, scale, dw, dh

    def postprocess(self, output, scale, dw, dh, img_shape):
        """后处理: NMS"""
        # YOLOv8输出shape: [1, 84, 8400] (COCO 80类)
        # 转为 [8400, 84]
        predictions = output[0].T

        # 获取置信度最高的类别
        class_scores = predictions[:, 4:]  # [8400, 80]
        max_scores = class_scores.max(axis=1)

        # 过滤低置信度
        mask = max_scores > self.conf_thres
        predictions = predictions[mask]
        max_scores = max_scores[mask]
        class_ids = class_scores[mask].argmax(axis=1)

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

        # 提取bbox (cx, cy, w, h → x1, y1, x2, y2)
        boxes = predictions[:, :4]
        x1 = boxes[:, 0] - boxes[:, 2] / 2
        y1 = boxes[:, 1] - boxes[:, 3] / 2
        x2 = boxes[:, 0] + boxes[:, 2] / 2
        y2 = boxes[:, 1] + boxes[:, 3] / 2

        # 还原到原图坐标
        x1 = (x1 - dw) / scale
        y1 = (y1 - dh) / scale
        x2 = (x2 - dw) / scale
        y2 = (y2 - dh) / scale

        # clip
        x1 = np.clip(x1, 0, img_shape[1])
        y1 = np.clip(y1, 0, img_shape[0])
        x2 = np.clip(x2, 0, img_shape[1])
        y2 = np.clip(y2, 0, img_shape[0])

        boxes_xyxy = np.stack([x1, y1, x2, y2], axis=1)

        # NMS
        indices = self._nms(boxes_xyxy, max_scores, self.iou_thres)

        return boxes_xyxy[indices], max_scores[indices], class_ids[indices]

    def _nms(self, boxes, scores, iou_threshold):
        """非极大值抑制"""
        indices = scores.argsort()[::-1]
        keep = []

        while len(indices) > 0:
            current = indices[0]
            keep.append(current)

            if len(indices) == 1:
                break

            rest = indices[1:]
            ious = self._compute_iou(boxes[current], boxes[rest])
            indices = rest[ious < iou_threshold]

        return keep

    def _compute_iou(self, box, boxes):
        """计算IoU"""
        x1 = np.maximum(box[0], boxes[:, 0])
        y1 = np.maximum(box[1], boxes[:, 1])
        x2 = np.minimum(box[2], boxes[:, 2])
        y2 = np.minimum(box[3], boxes[:, 3])

        inter = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
        area_box = (box[2] - box[0]) * (box[3] - box[1])
        area_boxes = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])

        return inter / (area_box + area_boxes - inter + 1e-6)

    def detect(self, img):
        """端到端检测"""
        # 预处理
        blob, scale, dw, dh = self.preprocess(img)

        # TensorRT推理
        input_name = self.engine.get_tensor_name(0)
        output_name = self.engine.get_tensor_name(1)

        self.context.set_input_shape(input_name, blob.shape)

        d_input = cuda.mem_alloc(blob.nbytes)
        output_shape = self.context.get_tensor_shape(output_name)
        output = np.empty(output_shape, dtype=np.float32)
        d_output = cuda.mem_alloc(output.nbytes)

        cuda.memcpy_htod_async(d_input, blob, self.stream)
        self.context.set_tensor_address(input_name, int(d_input))
        self.context.set_tensor_address(output_name, int(d_output))
        self.context.execute_async_v3(self.stream.handle)
        cuda.memcpy_dtoh_async(output, d_output, self.stream)
        self.stream.synchronize()

        d_input.free()
        d_output.free()

        # 后处理
        boxes, scores, class_ids = self.postprocess(
            output, scale, dw, dh, img.shape)

        return boxes, scores, class_ids

# 使用示例
def main():
    detector = YOLOv8TRT("yolov8n.engine")

    img = cv2.imread("test.jpg")

    # 预热
    for _ in range(10):
        detector.detect(img)

    # 性能测试
    N = 100
    start = time.perf_counter()
    for _ in range(N):
        boxes, scores, class_ids = detector.detect(img)
    elapsed = (time.perf_counter() - start) / N
    fps = 1.0 / elapsed

    print(f"Latency: {elapsed*1000:.1f}ms, FPS: {fps:.0f}")
    print(f"Detections: {len(boxes)}")

    # 可视化
    COCO_NAMES = ["person", "bicycle", "car", ...]  # 80类
    for box, score, cls_id in zip(boxes, scores, class_ids):  # zip按位置配对
        x1, y1, x2, y2 = box.astype(int)
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        label = f"{COCO_NAMES[cls_id]}: {score:.2f}"
        cv2.putText(img, label, (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

    cv2.imwrite("result.jpg", img)

import onnxruntime as ort
import numpy as np

class ONNXInference:
    """ONNX Runtime多后端推理"""

    def __init__(self, model_path, device="gpu"):
        # 根据设备选择ExecutionProvider
        providers = self._get_providers(device)

        # Session选项
        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = (
            ort.GraphOptimizationLevel.ORT_ENABLE_ALL)
        sess_options.intra_op_num_threads = 4
        sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL

        self.session = ort.InferenceSession(
            model_path, sess_options, providers=providers)

        # 获取输入输出信息
        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [o.name for o in self.session.get_outputs()]

        print(f"Using: {self.session.get_providers()}")

    def _get_providers(self, device):
        if device == "gpu":
            return [
                ("TensorrtExecutionProvider", {
                    "trt_max_workspace_size": 2 << 30,
                    "trt_fp16_enable": True,
                    "trt_engine_cache_enable": True,
                    "trt_engine_cache_path": "./trt_cache",
                }),
                ("CUDAExecutionProvider", {
                    "device_id": 0,
                    "arena_extend_strategy": "kNextPowerOfTwo",
                }),
                "CPUExecutionProvider",
            ]
        elif device == "openvino":
            return [
                ("OpenVINOExecutionProvider", {
                    "device_type": "CPU_FP32",
                }),
                "CPUExecutionProvider",
            ]
        else:
            return ["CPUExecutionProvider"]

    def infer(self, input_data):
        return self.session.run(
            self.output_names, {self.input_name: input_data})

# 使用
infer = ONNXInference("resnet50.onnx", device="gpu")
output = infer.infer(np.random.randn(1, 3, 224, 224).astype(np.float32))
print(f"Output shape: {output[0].shape}")

import tritonclient.http as httpclient
import tritonclient.grpc as grpcclient
import numpy as np

class TritonClient:
    """Triton推理客户端"""

    def __init__(self, url="localhost:8000", model_name="resnet50",
                 protocol="http"):
        if protocol == "http":
            self.client = httpclient.InferenceServerClient(url=url)
        else:
            self.client = grpcclient.InferenceServerClient(url=url)

        self.model_name = model_name
        self.protocol = protocol

    def infer(self, input_data: np.ndarray):
        """发送推理请求"""
        if self.protocol == "http":
            inputs = [httpclient.InferInput(
                "input", input_data.shape, "FP32")]
            inputs[0].set_data_from_numpy(input_data)

            outputs = [httpclient.InferRequestedOutput("output")]

            result = self.client.infer(
                self.model_name, inputs, outputs=outputs)
            return result.as_numpy("output")
        else:
            inputs = [grpcclient.InferInput(
                "input", input_data.shape, "FP32")]
            inputs[0].set_data_from_numpy(input_data)

            outputs = [grpcclient.InferRequestedOutput("output")]

            result = self.client.infer(
                self.model_name, inputs, outputs=outputs)
            return result.as_numpy("output")

    def check_health(self):
        """检查服务状态"""
        return {
            "live": self.client.is_server_live(),
            "ready": self.client.is_server_ready(),
            "model_ready": self.client.is_model_ready(self.model_name),
        }

# 启动Triton Server
"""
docker run --gpus=all --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 \
    -v $(pwd)/model_repository:/models \
    nvcr.io/nvidia/tritonserver:24.01-py3 \
    tritonserver --model-repository=/models
"""

# 使用
client = TritonClient("localhost:8000", "resnet50")
print(client.check_health())

input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = client.infer(input_data)
print(f"Top-1: class {np.argmax(output)}")

from nvidia.dali import pipeline_def, fn, types
from nvidia.dali.plugin.pytorch import DALIGenericIterator

@pipeline_def(batch_size=32, num_threads=4, device_id=0)
def dali_pipeline():
    """DALI GPU加速预处理"""
    # 读取文件
    jpegs, labels = fn.readers.file(
        file_root="./images/", random_shuffle=True)

    # GPU解码（比CPU OpenCV快5-10x）
    images = fn.decoders.image(jpegs, device="mixed",
                                output_type=types.RGB)

    # GPU上做resize
    images = fn.resize(images, device="gpu",
                       resize_x=224, resize_y=224)

    # GPU上做归一化
    images = fn.crop_mirror_normalize(
        images,
        device="gpu",
        dtype=types.FLOAT,
        output_layout="CHW",
        mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],
        std=[0.229 * 255, 0.224 * 255, 0.225 * 255],
    )

    return images, labels

# 使用DALI替代PyTorch DataLoader
pipe = dali_pipeline()
pipe.build()
dali_iter = DALIGenericIterator(pipe, ["images", "labels"])

for batch in dali_iter:
    images = batch[0]["images"]  # 已在GPU上
    labels = batch[0]["labels"]
    # 直接送入模型，无需额外transfer

import threading
import queue
import time
import numpy as np

class InferencePipeline:
    """三级流水线: 预处理 | 推理 | 后处理"""

    def __init__(self, model, preprocess_fn, postprocess_fn,
                 batch_size=8, max_queue_size=100):
        self.model = model
        self.preprocess_fn = preprocess_fn
        self.postprocess_fn = postprocess_fn
        self.batch_size = batch_size

        # 三个队列连接三个阶段
        self.raw_queue = queue.Queue(max_queue_size)
        self.infer_queue = queue.Queue(max_queue_size)
        self.result_queue = queue.Queue(max_queue_size)

        self.running = False

    def start(self):
        self.running = True
        # 启动三个工作线程
        threading.Thread(target=self._preprocess_worker, daemon=True).start()
        threading.Thread(target=self._infer_worker, daemon=True).start()
        threading.Thread(target=self._postprocess_worker, daemon=True).start()

    def _preprocess_worker(self):
        """预处理线程(CPU)"""
        batch = []
        while self.running:
            try:  # try/except捕获异常
                item = self.raw_queue.get(timeout=0.01)
                batch.append(self.preprocess_fn(item))

                if len(batch) >= self.batch_size:
                    self.infer_queue.put(np.stack(batch))
                    batch = []
            except queue.Empty:
                if batch:  # 超时时发送不完整的batch
                    self.infer_queue.put(np.stack(batch))
                    batch = []

    def _infer_worker(self):
        """推理线程(GPU)"""
        while self.running:
            try:
                batch = self.infer_queue.get(timeout=0.1)
                output = self.model.infer(batch)
                self.result_queue.put(output)
            except queue.Empty:
                continue

    def _postprocess_worker(self):
        """后处理线程(CPU)"""
        while self.running:
            try:
                output = self.result_queue.get(timeout=0.1)
                results = self.postprocess_fn(output)
                # 回调或存储结果
            except queue.Empty:
                continue

    def submit(self, item):
        """提交推理请求"""
        self.raw_queue.put(item)

    def stop(self):
        self.running = False