模型边缘部署：移动端、IoT 与 MLOps Pipeline

from dataclasses import dataclass
from enum import Enum

class DeviceType(Enum):
    MOBILE = "mobile"
    IOT_GATEWAY = "iot_gateway"
    EMBEDDED = "embedded"
    AUTOMOTIVE = "automotive"

@dataclass
class DeviceConstraint:
    device_type: DeviceType
    max_model_size_mb: float
    max_latency_ms: float
    max_power_mw: float
    has_npu: bool

# 典型边缘设备约束参数
MOBILE_CONSTRAINT = DeviceConstraint(
    device_type=DeviceType.MOBILE,
    max_model_size_mb=50.0,
    max_latency_ms=100.0,
    max_power_mw=500.0,
    has_npu=True
)

def check_deployment_feasibility(
    model_size_mb: float,
    latency_ms: float,
    constraint: DeviceConstraint
) -> bool:
    """检查模型是否满足目标设备的部署约束"""
    return (
        model_size_mb <= constraint.max_model_size_mb
        and latency_ms <= constraint.max_latency_ms
    )

print(check_deployment_feasibility(35, 80, MOBILE_CONSTRAINT))

import platform
import psutil
import subprocess

def detect_edge_hardware():
    """检测边缘设备的硬件能力，用于推理引擎选型"""
    info = {}

    # CPU 信息
    info["cpu_cores"] = psutil.cpu_count(logical=False)
    info["cpu_freq_mhz"] = psutil.cpu_freq().current if psutil.cpu_freq() else 0

    # 内存
    info["ram_mb"] = psutil.virtual_memory().total / (1024 ** 2)

    # GPU 加速（macOS 检查 Metal）
    if platform.system() == "Darwin":
        result = subprocess.run(
            ["system_profiler", "SPDisplaysDataType"],
            capture_output=True, text=True
        )
        info["has_gpu"] = "Apple M" in result.stdout

    # 平台判定
    if platform.machine() in ("arm64", "aarch64"):
        info["arch"] = "ARM"
        info["has_npu"] = platform.system() == "Darwin"  # Apple Neural Engine
    else:
        info["arch"] = "x86_64"
        info["has_npu"] = False

    return info

print(detect_edge_hardware())

import torch
import torch.nn as nn

# PTQ: 训练后静态量化（推荐用于边缘部署）
def apply_ptq(model: nn.Module, calib_dataloader, num_batches=100):
    """对训练好的模型执行训练后量化"""
    model.eval()

    # 融合 Conv + BN 层（量化前必须）
    model = torch.ao.quantization.fuse_modules(
        model, [["conv1", "bn1"], ["conv2", "bn2"]], inplace=True
    )

    # 配置量化参数
    model.qconfig = torch.ao.quantization.get_default_qconfig("qnnpack")
    torch.ao.quantization.prepare(model, inplace=True)

    # 校准：用少量数据激活统计信息收集
    with torch.no_grad():
        for i, (data, _) in enumerate(calib_dataloader):
            if i >= num_batches:
                break
            model(data)

    # 转换为量化模型
    torch.ao.quantization.convert(model, inplace=True)
    return model

# 验证量化效果
print(f"Model size reduced: {get_model_size(quantized_model):.1f} MB")

import torch
import torch.nn as nn
import torch.quantization as tq

# QAT: 量化感知训练
def apply_qat(model: nn.Module, train_dataloader, epochs=5):
    """量化感知训练流程"""
    model.eval()

    # 先融合模块
    model = torch.ao.quantization.fuse_modules(
        model, [["conv1", "bn1"], ["conv2", "bn2"]], inplace=True
    )

    # 配置 QAT
    model.qconfig = torch.ao.quantization.get_default_qat_qconfig("qnnpack")
    torch.ao.quantization.prepare_qat(model, inplace=True)

    # 训练过程中模拟量化
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(epochs):
        model.train()
        for data, target in train_dataloader:
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

    # 转换为最终量化模型
    model.eval()
    torch.ao.quantization.convert(model, inplace=True)
    return model

import torch
import torch.onnx

# 从 PyTorch 导出 ONNX 模型
def export_to_onnx(model, sample_input, output_path, opset_version=17):
    """将 PyTorch 模型导出为 ONNX 格式"""
    model.eval()

    torch.onnx.export(
        model,
        sample_input,
        output_path,
        export_params=True,
        opset_version=opset_version,
        do_constant_folding=True,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={
            "input": {0: "batch_size"},
            "output": {0: "batch_size"}
        }
    )
    print(f"Exported to {output_path}")

    # 验证导出模型
    import onnx
    onnx_model = onnx.load(output_path)
    onnx.checker.check_model(onnx_model)
    print("ONNX model validated")

    # 打印模型信息
    print(f"Opset version: {onnx_model.opset_import[0].version}")
    print(f"Num nodes: {len(onnx_model.graph.node)}")

import onnxruntime as ort
import numpy as np

# 使用 ONNX Runtime 进行推理
class ONNXInferencer:
    def __init__(self, model_path):
        # 选择执行提供者（优先级从高到低）
        providers = [
            ("CUDAExecutionProvider", {}),
            ("CoreMLExecutionProvider", {}),
            ("CPUExecutionProvider", {})
        ]
        self.session = ort.InferenceSession(
            model_path,
            providers=providers
        )
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name

        # 获取输入形状信息
        shape = self.session.get_inputs()[0].shape
        print(f"Expected input shape: {shape}")

    def predict(self, input_data: np.ndarray) -> np.ndarray:
        result = self.session.run(
            [self.output_name],
            {self.input_name: input_data}
        )
        return result[0]

# 性能测试
inferencer = ONNXInferencer("model.onnx")
import time
start = time.perf_counter()
for _ in range(100):
    inferencer.predict(np.random.randn(1, 3, 224, 224).astype(np.float32))
print(f"Average latency: {(time.perf_counter() - start) / 100 * 1000:.1f} ms")

import coremltools as ct
import torch

# 将 PyTorch 模型转换为 CoreML 格式
def convert_to_coreml(model, sample_input, output_path):
    """PyTorch -> TorchScript -> CoreML"""
    model.eval()

    # 1. 导出为 TorchScript
    traced_model = torch.jit.trace(model, sample_input)

    # 2. 转换为 ML Program（CoreML 新格式）
    mlmodel = ct.convert(
        traced_model,
        inputs=[ct.TensorType(shape=sample_input.shape)],
        convert_to="mlprogram",  # iOS 15+ 推荐格式
        compute_precision=ct.precision.FLOAT16,
        minimum_deployment_target=ct.target.iOS15
    )

    # 3. 添加元数据
    mlmodel.short_description = "Image classification model"
    mlmodel.author = "AI Master Site"
    mlmodel.license = "MIT"

    # 4. 保存
    mlmodel.save(output_path)
    print(f"Saved CoreML model to {output_path}")

    # 5. 验证模型
    import coremltools
    spec = coremltools.utils.load_spec(output_path)
    print(f"Input: {spec.description.input}")
    print(f"Output: {spec.description.output}")

import tensorflow as tf

# 将模型转换为 TFLite 并启用 NNAPI 委托
def convert_to_tflite(model, output_path):
    """转换模型为 TFLite 格式并优化移动端部署"""
    converter = tf.lite.TFLiteConverter.from_keras_model(model)

    # 启用优化
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.target_spec.supported_types = [tf.float16]

    # 导出 TFLite 模型
    tflite_model = converter.convert()
    with open(output_path, "wb") as f:
        f.write(tflite_model)
    print(f"Saved TFLite model: {len(tflite_model) / 1024:.1f} KB")

# 在 Android 上运行 TFLite（Java/Kotlin 侧模拟）
tflite_code = """
// Android 侧加载和运行 TFLite 模型
val interpreter = Interpreter(loadModelFile(context, "model.tflite"))

// 配置 NNAPI 委托
val options = Interpreter.Options()
options.setUseNNAPI(true)  // 启用 Android NPU 加速
options.setNumThreads(4)

val inputBuffer = TensorBuffer.createFixedSize(intArrayOf(1, 224, 224, 3), DataType.FLOAT32)
val outputBuffer = TensorBuffer.createFixedSize(intArrayOf(1, 1000), DataType.FLOAT32)

inputBuffer.loadImage(image)
interpreter.run(inputBuffer.buffer, outputBuffer.buffer)
val predictions = outputBuffer.floatArray
"""
print(tflite_code)

import tensorrt as trt
import numpy as np

# 使用 TensorRT Python API 构建推理引擎
def build_engine(onnx_path, engine_path):
    """将 ONNX 模型编译为 TensorRT Engine"""
    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

    # 创建 builder 和 network
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    )

    # 解析 ONNX 模型
    parser = trt.OnnxParser(network, TRT_LOGGER)
    with open(onnx_path, "rb") as f:
        parser.parse(f.read())
    print(f"Parsed ONNX: {network.num_layers} layers")

    # 配置构建参数
    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 30  # 1 GB 工作空间
    config.set_flag(trt.BuilderFlag.FP16)  # 启用 FP16 推理

    # 构建引擎
    engine = builder.build_serialized_network(network, config)
    with open(engine_path, "wb") as f:
        f.write(engine)
    print(f"Built engine: {len(engine) / 1024:.1f} KB")
    return engine

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

class TensorRTInferencer:
    def __init__(self, engine_path):
        TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, "rb") as f:
            self.runtime = trt.Runtime(TRT_LOGGER)
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        print(f"Engine loaded: {self.engine.num_io_tensors} tensors")

    def infer(self, input_data: np.ndarray) -> np.ndarray:
        """执行单次推理"""
        # 分配内存
        h_input = cuda.pagelocked_empty(
            trt.volume(self.engine.get_tensor_shape(0)), dtype=np.float32
        )
        h_output = cuda.pagelocked_empty(
            trt.volume(self.engine.get_tensor_shape(1)), dtype=np.float32
        )
        d_input = cuda.mem_alloc(h_input.nbytes)
        d_output = cuda.mem_alloc(h_output.nbytes)

        # 绑定地址
        self.context.set_tensor_address("input", int(d_input))
        self.context.set_tensor_address("output", int(d_output))

        # 执行推理
        np.copyto(h_input, input_data.ravel())
        stream = cuda.Stream()
        cuda.memcpy_htod_async(d_input, h_input, stream)
        self.context.execute_async_v3(stream_handle=stream.handle)
        cuda.memcpy_dtoh_async(h_output, d_output, stream)
        stream.synchronize()

        return h_output

# 性能基准测试
inferencer = TensorRTInferencer("model_fp16.engine")
times = []
for _ in range(200):
    start = time.perf_counter()
    inferencer.infer(np.random.randn(1, 3, 224, 224).astype(np.float32))
    times.append((time.perf_counter() - start) * 1000)
print(f"P50: {np.percentile(times, 50):.1f}ms, P99: {np.percentile(times, 99):.1f}ms")

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

# 结构化剪枝示例：移除不重要的卷积通道
def apply_structured_pruning(model, amount=0.3):
    """对模型中的 Conv2d 层执行 L1 结构化剪枝"""
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Conv2d):
            # L1 非结构化剪枝
            prune.l1_unstructured(module, name="weight", amount=amount)
            # 将剪枝永久化
            prune.remove(module, "weight")

    return model

# 计算剪枝后的稀疏度
def calc_sparsity(model):
    total = 0
    zeros = 0
    for name, param in model.named_parameters():
        if "weight" in name:
            total += param.numel()
            zeros += (param == 0).sum().item()
    return zeros / total * 100

pruned_model = apply_structured_pruning(model, amount=0.4)
print(f"Sparsity after pruning: {calc_sparsity(pruned_model):.1f}%")

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    """知识蒸馏损失：结合硬标签和软标签"""
    def __init__(self, temperature=4.0, alpha=0.5):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha  # 蒸馏损失权重
        self.hard_loss = nn.CrossEntropyLoss()

    def forward(self, student_logits, teacher_logits, labels):
        # 硬标签损失
        hard = self.hard_loss(student_logits, labels)

        # 软标签蒸馏损失 (KL 散度)
        soft_student = F.log_softmax(
            student_logits / self.temperature, dim=1
        )
        soft_teacher = F.softmax(
            teacher_logits / self.temperature, dim=1
        )
        soft = F.kl_div(
            soft_student, soft_teacher, reduction="batchmean"
        ) * (self.temperature ** 2)

        return self.alpha * soft + (1 - self.alpha) * hard

# 蒸馏训练循环
teacher_model.eval()
for data, labels in train_loader:
    with torch.no_grad():
        teacher_logits = teacher_model(data)
    student_logits = student_model(data)
    loss = distill_loss(student_logits, teacher_logits, labels)
    loss.backward()
    optimizer.step()

import torch
import torch.onnx
import onnx
from onnxsim import simplify
import coremltools as ct

def pytorch_to_coreml(model, input_shape=(1, 3, 224, 224),
                      onnx_path="model.onnx",
                      mlmodel_path="model.mlpackage"):
    """完整转换流程：PyTorch -> ONNX -> CoreML"""
    model.eval()
    dummy_input = torch.randn(*input_shape)

    # Step 1: PyTorch -> ONNX
    print("Step 1: Exporting to ONNX...")
    torch.onnx.export(
        model, dummy_input, onnx_path,
        opset_version=17,
        input_names=["image"],
        output_names=["predictions"],
        dynamic_axes={
            "image": {0: "batch"},
            "predictions": {0: "batch"}
        }
    )

    # Step 2: Simplify ONNX
    print("Step 2: Simplifying ONNX...")
    onnx_model = onnx.load(onnx_path)
    simplified, check = simplify(onnx_model)
    onnx.save(simplified, "model_simplified.onnx")
    print(f"  Nodes: {len(onnx_model.graph.node)} -> {len(simplified.graph.node)}")

    # Step 3: ONNX -> CoreML
    print("Step 3: Converting to CoreML...")
    mlmodel = ct.convert(
        "model_simplified.onnx",
        convert_to="mlprogram",
        compute_precision=ct.precision.FLOAT16,
        minimum_deployment_target=ct.target.iOS15
    )

    # 添加图像输入类型
    mlmodel.input_description["image"] = "Input image (224x224 RGB)"
    mlmodel.output_description["predictions"] = "Class probabilities"
    mlmodel.save(mlmodel_path)

    # Step 4: 验证精度
    print("Step 4: Verifying CoreML model...")
    coreml_model = ct.models.MLModel(mlmodel_path)
    coreml_input = {"image": dummy_input.numpy()}
    coreml_output = coreml_model.predict(coreml_input)

    with torch.no_grad():
        pytorch_output = model(dummy_input).numpy()

    diff = abs(coreml_output["predictions"] - pytorch_output).max()
    print(f"  Max difference: {diff:.6f}")
    print(f"  Model saved to: {mlmodel_path}")
    return mlmodel_path

# 执行完整流程
# model = torch.hub.load("NVIDIA/DeepLearningExamples", "efficientnet_b0")
# pytorch_to_coreml(model)

import UIKit
import Vision
import CoreML

class ImageClassifier {
    private let model: VNCoreMLModel
    private let request: VNCoreMLRequest

    init() throws {
        // 加载 CoreML 模型
        guard let mlModel = try? ImageClassifierModel(
            configuration: MLModelConfiguration()
        ) else {
            fatalError("Failed to load model")
        }
        model = try VNCoreMLModel(for: mlModel.model)

        // 配置 Vision 请求
        request = VNCoreMLRequest(model: model) { [weak self] request, error in
            guard let results = request.results as? [VNClassificationObservation],
                  let top = results.first else { return }
            print("Classification: \(top.identifier) (\(top.confidence))")
        }

        // 预处理配置
        request.imageCropAndScaleOption = .centerCrop
    }

    func classify(image: UIImage) {
        guard let cgImage = image.cgImage else { return }
        let handler = VNImageRequestHandler(cgImage: cgImage)
        try? handler.perform([request])
    }

    func classifyBatch(images: [UIImage]) {
        // 批量推理：CoreML 支持 batch prediction
        let requests = images.map { image -> VNCoreMLRequest in
            let req = VNCoreMLRequest(model: model)
            req.imageCropAndScaleOption = .centerCrop
            return req
        }
        let handler = VNImageRequestHandler(cgImage: images[0].cgImage!)
        try? handler.perform(requests)
    }
}