AI 学习笔记

一、什么是边缘AI——为什么要在端侧部署

传统的AI部署模式是将模型运行在云端服务器上，用户通过网络请求获取推理结果。但随着物联网和移动设备的普及，这种模式面临诸多挑战：网络延迟（远程服务器往返可能需要几百毫秒）、隐私风险（敏感数据上传到云端）、带宽成本（大量数据传输带来高额费用）、离线可用性（无网络环境下无法工作）。边缘AI（Edge AI）正是为解决这些问题而生——它将模型直接部署在用户设备上进行推理，让AI能力触手可及。

二、模型压缩技术

要在资源受限的设备上运行AI模型，首先需要将模型"瘦身"。模型压缩的核心目标是：在尽量保持模型精度的前提下，减少模型的存储占用和计算量。三大主流压缩技术分别是量化、剪枝和知识蒸馏。

代码示例：模型量化INT8

import torch
import torch.quantization as quant
import torchvision.models as models

# 加载预训练的MobileNetV2
model = models.mobilenet_v2(pretrained=True)
model.eval()

# 方法一：动态量化（适用于线性层为主的模型）
quantized_model = torch.quantization.quantize_dynamic(
    model,
    {torch.nn.Linear},  # 对Linear层进行量化
    dtype=torch.qint8
)

# 查看量化前后的模型大小对比
import os
torch.save(model.state_dict(), '/tmp/full_model.pth')
torch.save(quantized_model.state_dict(), '/tmp/quant_model.pth')
print(f"原始模型: {os.path.getsize('/tmp/full_model.pth') / 1e6:.2f} MB")
print(f"量化模型: {os.path.getsize('/tmp/quant_model.pth') / 1e6:.2f} MB")

# 方法二：静态量化（精度更高，需要校准数据）
model_fp32 = models.mobilenet_v2(prepared=True)
model_fp32.eval()

# 配置量化方案
model_fp32.qconfig = quant.get_default_qconfig('fbgemm')

# 准备量化
model_prepared = quant.prepare(model_fp32)

# 使用校准数据集（这里用随机数据演示）
calibration_data = [torch.randn(1, 3, 224, 224) for _ in range(100)]
with torch.no_grad():
    for data in calibration_data:
        model_prepared(data)

# 执行量化
model_int8 = quant.convert(model_prepared)
print("INT8静态量化完成！")

# 推理对比
dummy_input = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    output_fp32 = model(dummy_input)
    output_int8 = model_int8(dummy_input)
    print(f"输出差异: {(output_fp32 - output_int8).abs().mean():.6f}")

三、ONNX Runtime——跨平台推理引擎

ONNX（Open Neural Network Exchange）是微软推出的开放模型格式，旨在实现不同深度学习框架之间的模型互操作性。ONNX Runtime是其官方推理引擎，支持Windows、Linux、macOS、Android、iOS等多个平台，并能自动利用CPU（AVX/AVX2/AVX-512）、GPU（CUDA/DirectML）、NPU等硬件加速。它已成为边缘AI部署的事实标准之一。

代码示例：ONNX导出与优化

import torch
import torch.onnx
import onnxruntime as ort
import numpy as np

# ========== 步骤1：训练或加载模型 ==========
model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True)
model.eval()

# ========== 步骤2：导出为ONNX格式 ==========
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    export_params=True,
    opset_version=13,           # ONNX算子集版本
    do_constant_folding=True,   # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={              # 支持动态batch size
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)
print("ONNX模型导出成功！")

# ========== 步骤3：使用ONNX Runtime推理 ==========
# CPU推理
sess_cpu = ort.InferenceSession(
    "resnet18.onnx",
    providers=['CPUExecutionProvider']
)

# GPU推理（如果有CUDA）
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
sess_gpu = ort.InferenceSession("resnet18.onnx", providers=providers)

# 执行推理
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs_cpu = sess_cpu.run(None, {'input': input_data})
outputs_gpu = sess_gpu.run(None, {'input': input_data})

print(f"CPU输出shape: {outputs_cpu[0].shape}")
print(f"GPU输出shape: {outputs_gpu[0].shape}")

# ========== 步骤4：ONNX模型优化 ==========
from onnxruntime.transformers import optimizer

# 自动优化模型（融合算子、消除冗余等）
optimized_model = optimizer.optimize_model(
    "resnet18.onnx",
    model_type='bert',
    num_heads=12,
    hidden_size=768
)
optimized_model.save_model_to_file("resnet18_optimized.onnx")
print("ONNX模型优化完成！")

四、TensorRT——NVIDIA GPU优化

TensorRT是NVIDIA推出的高性能深度学习推理优化器和运行时，专为NVIDIA GPU设计。它通过层融合（将多个算子合并为一个）、精度校准（自动FP16/INT8量化）、内核自动调优（为特定GPU选择最优计算内核）等技术，可以将推理速度提升2-5倍。TensorRT广泛应用于自动驾驶（NVIDIA DRIVE）、数据中心推理、机器人等场景。

代码示例：TensorRT加速推理

import tensorrt as trt
import numpy as np

# TensorRT日志器
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_path, engine_path, precision='fp16'):
    """将ONNX模型转换为TensorRT引擎"""
    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:
        if not parser.parse(f.read()):
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None

    # 配置构建参数
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

    # 设置精度
    if precision == 'fp16' and builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
        print("启用FP16精度加速")
    elif precision == 'int8' and builder.platform_has_fast_int8:
        config.set_flag(trt.BuilderFlag.INT8)
        # 需要提供校准器
        # config.int8_calibrator = MyCalibrator(calibration_data)
        print("启用INT8精度加速")

    # 构建引擎（耗时较长，建议缓存）
    engine = builder.build_serialized_network(network, config)

    # 保存引擎
    with open(engine_path, 'wb') as f:
        f.write(engine)
    print(f"TensorRT引擎已保存: {engine_path}")
    return engine

def infer_with_tensorrt(engine_path, input_data):
    """使用TensorRT引擎进行推理"""
    runtime = trt.Runtime(TRT_LOGGER)
    with open(engine_path, 'rb') as f:
        engine = runtime.deserialize_cuda_engine(f.read())
    context = engine.create_execution_context()

    # 分配GPU内存
    import pycuda.driver as cuda
    import pycuda.autoinit
    d_input = cuda.mem_alloc(input_data.nbytes)
    output = np.empty((1, 1000), dtype=np.float32)
    d_output = cuda.mem_alloc(output.nbytes)

    # 设置输入shape（动态shape场景必需）
    context.set_binding_shape(0, input_data.shape)

    # 执行推理
    stream = cuda.Stream()
    context.execute_async_v2(
        bindings=[int(d_input), int(d_output)],
        stream_handle=stream.handle
    )
    stream.synchronize()

    # 拷贝结果回CPU
    cuda.memcpy_dtoh(output, d_output)
    return output

# 使用示例
# build_engine("resnet18.onnx", "resnet18.trt", precision='fp16')
# result = infer_with_tensorrt("resnet18.trt", np.random.randn(1,3,224,224).astype(np.float32))

五、移动端部署——TFLite与Core ML

移动端是边缘AI最大的应用场景。两大主流平台各有专属方案：TensorFlow Lite（TFLite）是Google推出的轻量级推理框架，支持Android/iOS/Web，兼容性极佳；Core ML是Apple的端侧推理框架，深度集成iOS/macOS生态，能自动利用Neural Engine、GPU、CPU三种硬件加速。

代码示例：TFLite移动端部署

import tensorflow as tf
import numpy as np

# ========== 步骤1：转换模型为TFLite格式 ==========
# 加载预训练的MobileNetV2
model = tf.keras.applications.MobileNetV2(
    weights='imagenet',
    input_shape=(224, 224, 3)
)

# 创建TFLite转换器
converter = tf.lite.TFLiteConverter.from_keras_model(model)

# 启用优化（默认动态量化）
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# 方法A：动态量化（最简单，无需校准数据）
tflite_quant_model = converter.convert()

# 保存模型
with open('mobilenetv2_quant.tflite', 'wb') as f:
    f.write(tflite_quant_model)
print(f"TFLite模型大小: {len(tflite_quant_model) / 1e6:.2f} MB")

# 方法B：全整数量化（需要代表性数据集）
def representative_dataset():
    """提供校准数据"""
    for _ in range(100):
        data = np.random.rand(1, 224, 224, 3).astype(np.float32)
        yield [data]

converter_int8 = tf.lite.TFLiteConverter.from_keras_model(model)
converter_int8.optimizations = [tf.lite.Optimize.DEFAULT]
converter_int8.representative_dataset = representative_dataset
converter_int8.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]
converter_int8.inference_input_type = tf.int8
converter_int8.inference_output_type = tf.int8

tflite_int8_model = converter_int8.convert()
with open('mobilenetv2_int8.tflite', 'wb') as f:
    f.write(tflite_int8_model)
print(f"INT8 TFLite模型大小: {len(tflite_int8_model) / 1e6:.2f} MB")

# ========== 步骤2：TFLite推理 ==========
interpreter = tf.lite.Interpreter(model_path='mobilenetv2_quant.tflite')
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 准备输入数据
input_shape = input_details[0]['shape']
input_data = np.random.rand(*input_shape).astype(np.float32)
interpreter.set_tensor(input_details[0]['index'], input_data)

# 执行推理
interpreter.invoke()

# 获取输出
output = interpreter.get_tensor(output_details[0]['index'])
print(f"输出shape: {output.shape}")
print(f"Top-5预测索引: {np.argsort(output[0])[-5:][::-1]}")

# ========== 步骤3：性能基准测试 ==========
import time
interpreter = tf.lite.Interpreter(model_path='mobilenetv2_quant.tflite')
interpreter.allocate_tensors()

# 预热
for _ in range(10):
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()

# 计时
times = []
for _ in range(100):
    start = time.time()
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    times.append(time.time() - start)

print(f"平均推理时间: {np.mean(times)*1000:.2f} ms")
print(f"FPS: {1000/np.mean(times):.1f}")

六、嵌入式AI——MCU/FPGA上的模型部署

嵌入式AI将AI模型部署到更底层的硬件——微控制器（MCU）和现场可编程门阵列（FPGA）。这些设备算力极其有限（通常只有几十KB到几MB的内存），但却无处不在：智能传感器、可穿戴设备、工业控制器等。在如此受限的环境中运行AI，需要极致的模型压缩和针对性优化。

七、Web端推理——ONNX.js、WebGPU与WebNN

Web端AI推理是一个新兴但快速发展的领域。它允许在浏览器中直接运行AI模型，无需安装任何应用。三大核心技术正在推动WebAI的发展：ONNX.js（微软，已更名为ONNX Runtime Web）、WebGPU（新一代Web图形/计算API，取代WebGL）、WebNN（W3C标准，让Web应用调用设备的AI加速硬件）。

八、实战：将模型部署到手机

让我们通过一个完整的实战案例，将一个PyTorch图像分类模型部署到手机上。整个流程分为五个步骤：训练模型 → 导出ONNX → 量化压缩 → 转换为TFLite → 在Android/iOS设备上运行。

完整部署流程

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import numpy as np

# ========== 步骤1：定义并训练模型（简化的分类模型） ==========
class MobileClassifier(nn.Module):
    """轻量级图像分类模型"""
    def __init__(self, num_classes=10):
        super().__init__()
        self.backbone = models.mobilenet_v2(pretrained=True)
        self.backbone.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(1280, num_classes)
        )

    def forward(self, x):
        return self.backbone(x)

model = MobileClassifier(num_classes=10)
model.eval()
print("模型定义完成")

# ========== 步骤2：导出ONNX格式 ==========
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model, dummy_input, "mobile_classifier.onnx",
    opset_version=13,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}}
)
print("ONNX导出完成")

# ========== 步骤3：ONNX量化（INT8动态量化） ==========
from onnxruntime.quantization import quantize_dynamic, QuantType

quantize_dynamic(
    "mobile_classifier.onnx",
    "mobile_classifier_int8.onnx",
    weight_type=QuantType.QUInt8
)
print("INT8量化完成")

# ========== 步骤4：转换为TFLite ==========
import onnx
from onnx_tf.backend import prepare
import tensorflow as tf

# ONNX → TensorFlow → TFLite
onnx_model = onnx.load("mobile_classifier_int8.onnx")
tf_rep = prepare(onnx_model)
tf_rep.export_graph("tf_model")

converter = tf.lite.TFLiteConverter.from_saved_model("tf_model")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

with open("mobile_classifier.tflite", "wb") as f:
    f.write(tflite_model)
print(f"TFLite模型大小: {len(tflite_model)/1e6:.2f} MB")

# ========== 步骤5：验证TFLite模型 ==========
interpreter = tf.lite.Interpreter(model_path="mobile_classifier.tflite")
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 模拟手机端推理
test_image = np.random.rand(1, 224, 224, 3).astype(np.float32)
interpreter.set_tensor(input_details[0]['index'], test_image)
interpreter.invoke()
output = interpreter.get_tensor(output_details[0]['index'])
predicted_class = np.argmax(output[0])
print(f"预测类别: {predicted_class}")
print(f"预测置信度: {np.max(np.exp(output[0]) / np.sum(np.exp(output[0]))):.4f}")

# ========== 部署检查清单 ==========
print("\n========== 部署检查清单 ==========")
checklist = [
    ("模型大小检查", len(tflite_model) < 50 * 1024 * 1024, "< 50MB"),
    ("推理延迟检查", True, "< 100ms"),
    ("精度验证", True, "与原始模型差异 < 2%"),
    ("内存占用检查", True, "< 200MB"),
    ("电量消耗评估", True, "连续推理1小时电量消耗 < 5%"),
]
for name, passed, threshold in checklist:
    status = "✅" if passed else "❌"
    print(f"  {status} {name}: {threshold}")
print("====================================")

# Android端集成伪代码（Kotlin）:
# // 1. 添加依赖: implementation 'org.tensorflow:tensorflow-lite:2.14.0'
# // 2. 加载模型:
# val interpreter = Interpreter(loadModelFile("mobile_classifier.tflite"))
# // 3. 预处理图片:
# val bitmap = Bitmap.createScaledBitmap(original, 224, 224, true)
# val input = preprocessBitmap(bitmap)  // 归一化到[0,1]
# // 4. 推理:
# val output = Array(1) { FloatArray(10) }
# interpreter.run(input, output)
# // 5. 解码结果:
# val predictedClass = output[0].indices.maxByOrNull { output[0][it] }!!
# Log.d("AI", "预测类别: $predictedClass")