Input Clip (T=8, H=224, W=224, C=3)
        │
        ▼
┌───────────────────────┐
│  Swin-Tiny Backbone   │  ← ImageNet-21K 预训练，窗口注意力降低计算量
│  (patch_embed + 4 stages)    输出特征图 1/32 分辨率
└───────────────────────┘
        │
        ▼
┌───────────────────────┐
│  Fragment Sampling    │  ← 空间分片采样（8×8 grid），保留局部细节
│  + Temporal Pooling   │     时间维度 Attention Pooling (num_heads=4)
└───────────────────────┘
        │
        ▼
┌───────────────────────┐
│  Quality Regressor    │  ← 2-layer MLP (512→128→1) + Dropout(0.1)
└───────────────────────┘
        │
        ▼
    MOS Score ∈ [1,5]

# Stage-1: 预训练冻结骨干，仅训练回归头 + Fragment Sampling 模块
# Stage-2: 解冻 Stage-3/4，小倍率微调全模型

graph LR
A[PyTorch Checkpoint] --> B[torch.onnx.export opset=17]
B --> C[ONNX Graph Surgeon]
C --> D[Fold Constant / Eliminate Dead Nodes]
D --> E[ONNX Runtime / TensorRT Engine Build]
E --> F[INT8 Calibration (Entropy/Percentile)]
F --> G[Profile & Validate]

def plcc_loss(pred: torch.Tensor, target: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    """
    最大化 PLCC <=> 最小化 1 - ρ
    梯度特性：当 pred 与 target 线性相关时梯度为 0；非线性畸变区域梯度大，强制拉直拟合曲线。
    """
    pred_mean = pred.mean(dim=0, keepdim=True)
    target_mean = target.mean(dim=0, keepdim=True)
    
    v_pred = pred - pred_mean
    v_target = target - target_mean
    
    cov = (v_pred * v_target).sum(dim=0)
    std_pred = torch.sqrt((v_pred ** 2).sum(dim=0) + eps)
    std_target = torch.sqrt((v_target ** 2).sum(dim=0) + eps)
    
    rho = cov / (std_pred * std_target + eps)
    return 1.0 - rho.mean()  # scalar loss

def rank_loss(pred: torch.Tensor, target: torch.Tensor, margin: float = 0.1) -> torch.Tensor:
    """
    仅对「序关系不一致」的样本对施加惩罚，保序不保值。
    复杂度 O(N^2) -> 采用「邻序采样」降为 O(N log N)：
    1. 按 target 排序得索引 idx_sorted
    2. 仅构造相邻 K=3 个正序对 + 随机采样 M=16 个负序对
    """
    with torch.no_grad():
        idx_sorted = target.argsort(dim=0, descending=True)
        # 正序对
        pos_pairs = [(idx_sorted[i], idx_sorted[i+1]) for i in range(len(idx_sorted)-1)]
        # 负序对：随机打乱后配对
        neg_idx = idx_sorted[torch.randperm(len(idx_sorted))]
        neg_pairs = [(idx_sorted[i], neg_idx[i]) for i in range(len(idx_sorted))]
    
    loss = 0.0
    for i, j in pos_pairs + neg_pairs:
        sign = 1.0 if target[i] > target[j] else -1.0
        diff = pred[i] - pred[j]
        loss += F.relu(margin - sign * diff)
    return loss / (len(pos_pairs) + len(neg_pairs))

# 1. 导出 ONNX (动态 Batch, 固定 T=8, H=W=224)
python export_onnx.py --model fast_vqa_m --opset 17 --dynamic_batch --simplify

# 2. 准备校准集：覆盖全失真类型、全分辨率、全运动幅度的 500 片段
#    关键：校准集分布必须 ≈ 线上流量分布，否则 KL 散度最小化会偏移
python prepare_calib.py --src /data/online_sampled --dst /calib --num 500 --strategy stratified

# 3. TensorRT INT8 Engine Build (trtexec)
trtexec --onnx=model.onnx 
        --int8 --calib=calib.cache 
        --minShapes=input:1x8x3x224x224 
        --optShapes=input:4x8x3x224x224 
        --maxShapes=input:8x8x3x224x224 
        --workspace=4096 
        --saveEngine=model_int8.trt 
        --verbose 2>&1 | tee trt_build.log

# 4. ONNX Runtime INT8 (VNNI/AMX 加速)
python -m onnxruntime.quantization.quantize_static 
       --input_model model.onnx 
       --output_model model_ort_int8.onnx 
       --calibration_data_reader CalibReader('/calib') 
       --quant_format QOperator 
       --activation_type QUInt8 --weight_type QInt8 
       --optimize_model True

def verify_numerics(onnx_fp32, onnx_int8, trt_int8, calib_loader, atol=1e-2, rtol=1e-3):
    """
    逐层输出对齐 + 端到端指标对齐
    """
    for batch in calib_loader:
        out_fp32 = ort_fp32.run(batch)
        out_ort = ort_int8.run(batch)
        out_trt = trt_int8.run(batch)
        
        # 1. 端到端 MOS 分数
        assert np.allclose(out_fp32, out_ort, atol=atol, rtol=rtol), "ORT INT8 drift"
        assert np.allclose(out_fp32, out_trt, atol=atol, rtol=rtol), "TRT INT8 drift"
        
        # 2. 关键中间层 (Attn Output, Fragment Feat) 逐层 Hook 对比
        #    使用 ONNX Runtime SessionIOBinding 获取中间张量
        for layer_name in CRITICAL_LAYERS:
            fp32_feat = get_intermediate(ort_fp32, layer_name, batch)
            int8_feat = get_intermediate(ort_int8, layer_name, batch)
            cos_sim = cosine_similarity(fp32_feat.flatten(), int8_feat.flatten())
            assert cos_sim > 0.999, f"Layer {layer_name} cosine sim {cos_sim:.6f} < 0.999"
    print("✅ Numerical verification passed.")

# PyTorch FX Graph Mode QAT (支持自定义算子)
from torch.ao.quantization import QConfigMapping, get_default_qat_qconfig
from torch.ao.quantization.quantize_fx import prepare_qat_fx, convert_fx

qconfig = get_default_qat_qconfig('x86')  # 或 'qnnpack' for ARM
qconfig_mapping = QConfigMapping().set_global(qconfig)

# 关键：为 Swin Attention 注册自定义 QAT 模式
from torch.ao.quantization import Quantizer
class SwinAttentionQuantizer(Quantizer):
    @classmethod
    def get_supported_quantized_modules(cls):
        return [SwinAttention]  # 自定义模块
    
    def quantize(self, module, **kwargs):
        # 手动插入 FakeQuantize 于 Q/K/V 投影输出、Attn Output
        # 并将 Softmax 替换为 Log-Int-Softmax 近似实现
        return QuantizedSwinAttention(module)

model_qat = prepare_qat_fx(model_fp32, qconfig_mapping, example_inputs=dummy_input)
# 微调 5 Epochs, LR=1e-5
train_qat(model_qat, epochs=5, lr=1e-5)
model_int8 = convert_fx(model_qat.eval())
torch.jit.save(torch.jit.trace(model_int8, dummy_input), "model_qat_int8.pt")

name: "fast_vqa_m_int8"
backend: "tensorrt"
max_batch_size: 8
input [
  { name: "input" data_type: TYPE_FP32 dims: [8, 3, 224, 224] }
]
output [
  { name: "mos" data_type: TYPE_FP32 dims: [1] }
]
instance_group [
  { count: 2 kind: KIND_GPU gpus: [0, 1] }  # 单 GPU 2 实例，避免显存 OOM
  { count: 2 kind: KIND_GPU gpus: [2, 3] }
]
dynamic_batching {
  preferred_batch_size: [4, 8]
  max_queue_delay_microseconds: 500  # 关键：延迟预算内最大化批处理
  preserve_ordering: true
}
optimization {
  input_pinned_memory: true
  output_pinned_memory: true
  cuda_graphs: true  # 捕获 CUDA Graph 消除 Kernel Launch 开销
}
model_repository_agent {  # 热更新
  poll_seconds: 30
}

# client_sdk/vqa_client.py
import tritonclient.grpc.aio as grpcclient
import numpy as np
from collections import deque
import asyncio

class VQAClient:
    def __init__(self, url="localhost:8001", model_name="fast_vqa_m_int8", concurrency=32):
        self._client = grpcclient.InferenceServerClient(url=url, verbose=False)
        self._model_name = model_name
        self._sem = asyncio.Semaphore(concurrency)  # 客户端背压控制
        self._buffer_pool = deque([np.empty((8, 3, 224, 224), dtype=np.float32) for _ in range(concurrency * 2)])

    async def infer_async(self, frames: np.ndarray) -> float:
        """
        frames: (T, H, W, 3) uint8 RGB -> 内部预处理 -> (8, 3, 224, 224) fp32
        返回: MOS float
        """
        # 1. 从池获取 buffer，避免频繁 malloc
        buf = self._buffer_pool.popleft() if self._buffer_pool else np.empty((8,3,224,224), np.float32)
        preprocess(frames, buf)  # 原地操作，零拷贝
        
        inputs = [grpcclient.InferInput("input", buf.shape, "FP32")]
        inputs[0].set_data_from_numpy(buf, binary_data=True)
        outputs = [grpcclient.InferRequestedOutput("mos", binary_data=True)]
        
        async with self._sem:
            result = await self._client.async_infer(self._model_name, inputs, outputs=outputs)
        
        mos = result.as_numpy("mos")[0, 0]
        self._buffer_pool.append(buf)  # 归还 buffer
        return float(mos)

    async def stream_infer(self, frame_queue: asyncio.Queue, result_queue: asyncio.Queue):
        """流式推理入口：滑动窗口 T=8, stride=4"""
        window = []
        while True:
            frame = await frame_queue.get()
            window.append(frame)
            if len(window) == 8:
                mos = await self.infer_async(np.stack(window))
                await result_queue.put((window[0].timestamp, mos))  # 对齐首帧时间戳
                window = window[4:]  # stride=4 重叠

class MultiTaskHead(nn.Module):
    def __init__(self, in_dim=512, num_distortions=14):
        super().__init__()
        self.mos_head = nn.Sequential(
            nn.Linear(in_dim, 128), nn.GELU(), nn.Dropout(0.1), nn.Linear(128, 1)
        )
        self.dist_head = nn.Sequential(
            nn.Linear(in_dim, 128), nn.GELU(), nn.Linear(128, num_distortions)
        )
        # 共享特征提取层
        self.shared = nn.Identity()  # 直接复用 Backbone 最后一层 Fragment Feature

    def forward(self, x):
        feat = self.shared(x)  # [B, 512]
        mos = self.mos_head(feat).squeeze(-1)
        dist_logits = self.dist_head(feat)  # [B, 14]
        return mos, dist_logits

# 损失函数扩展
def multi_task_loss(mos_pred, mos_gt, dist_pred, dist_gt, lambda_dist=0.3):
    loss_mos = plcc_loss(mos_pred, mos_gt) + 0.2 * F.l1_loss(mos_pred, mos_gt)
    loss_dist = F.binary_cross_entropy_with_logits(dist_pred, dist_gt.float())  # 多标签
    return loss_mos + lambda_dist * loss_dist

# rules/root_cause.py
from dataclasses import dataclass
from enum import Enum

class DistortionType(Enum):
    FREEZE = "freeze_frame"
    BLOCK_ARTIFACT = "block_artifact"
    COLOR_CAST = "color_cast"
    BLUR = "blur"
    NOISE = "sensor_noise"
    # ... 共 14 类

@dataclass
class Context:
    mos: float
    dist_probs: Dict[DistortionType, float]  # sigmoid 后概率
    net_stats: Dict  # {loss_rate, rtt, jitter, bandwidth_kbps}
    codec_stats: Dict # {qp_avg, frame_type_ratio, decode_error_cnt}

def diagnose(ctx: Context) -> Dict[str, str]:
    """
    返回: {root_cause: "network|encoder|decoder|capture", action: "lower_bitrate|switch_codec|alert_user|..."}
    """
    # 规则 1: 网络主导
    if ctx.net_stats['loss_rate'] > 0.05 and ctx.dist_probs[DistortionType.FREEZE] > 0.7:
        return {"root_cause": "network", "action": "force_audio_only", "confidence": 0.95}
    
    # 规则 2: 编码器回归
    if ctx.codec_stats['qp_avg'] < 30 and ctx.dist_probs[DistortionType.BLOCK_ARTIFACT] > 0.6:
        return {"root_cause": "encoder", "action": "rollback_codec_version", "confidence": 0.9}
    
    # 规则 3: 采集端问题
    if ctx.dist_probs[DistortionType.NOISE] > 0.8 and ctx.net_stats['loss_rate'] < 0.01:
        return {"root_cause": "capture", "action": "notify_device_vendor", "confidence": 0.85}
    
    # 规则 4: 解码端异常
    if ctx.codec_stats['decode_error_cnt'] > 10 and ctx.dist_probs[DistortionType.BLOCK_ARTIFACT] > 0.5:
        return {"root_cause": "decoder", "action": "reset_decoder_pipeline", "confidence": 0.9}
    
    return {"root_cause": "unknown", "action": "collect_more_data", "confidence": 0.3}

graph TD
A[Teacher: Swin-Tiny + FragHead<br/>28M params, SROCC 0.851] --> B[Student: MobileNetV3-Small + TinyTrans<br/>1.8M params]
B --> C[Channel Pruning (L1 Norm, 50% FLOPs)]
C --> D[KD Training: L2(feat) + KL(logits) + Rank]
D --> E[INT8 QAT (PTQ failed on NPU)]
E --> F[RKNN/NCNN/TFLite Export]
F --> G[RK3588 NPU / ARM CPU Benchmark]

def distillation_loss(student_feat, teacher_feat, student_logits, teacher_logits, target_mos):
    # 1. 特征模仿 (Fragment Level)
    loss_feat = F.mse_loss(student_feat, teacher_feat.detach())
    
    # 2. Logits 蒸馏 (Temperature=2.0)
    T = 2.0
    loss_kl = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),
        F.softmax(teacher_logits / T, dim=-1),
        reduction='batchmean'
    ) * (T ** 2)
    
    # 3. 任务监督
    loss_task = plcc_loss(student_logits, target_mos)
    
    return loss_task + 0.5 * loss_feat + 0.3 * loss_kl

# 将模型输出原始分数映射到主观 MOS 尺度
def logistic5(x, beta1, beta2, beta3, beta4, beta5):
    """VQEG 标准 5 参数 Logistic 函数"""
    return beta1 * (0.5 - 1.0 / (1.0 + np.exp(beta2 * (x - beta3)))) + beta4 * x + beta5

from scipy.optimize import curve_fit
popt, _ = curve_fit(logistic5, model_raw_scores, subjective_mos, maxfev=10000)
mapped_scores = logistic5(model_raw_scores, *popt)

# 拟合后指标
plcc_mapped = pearsonr(mapped_scores, subjective_mos)[0]
srocc_mapped = spearmanr(mapped_scores, subjective_mos)[0]
rmse_mapped = np.sqrt(np.mean((mapped_scores - subjective_mos)**2))

# train.sh
export NCCL_DEBUG=INFO
export PYTHONFAULTHANDLER=1
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=$RANK --master_addr=$MASTER_ADDR 
    train.py 
    --cfg configs/fast_vqa_m_finetune.yaml 
    --data_root /mnt/data/vqa_domain 
    --output_dir /mnt/ckpt/fast_vqa_m_v2 
    --resume auto 
    --amp 
    --wandb_project "meeting-vqa" 
    2>&1 | tee -a logs/train_$(date +%Y%m%d_%H%M).log

# .gitlab-ci.yml
stages:
  - lint
  - unit_test
  - train_smoke
  - export_onnx
  - quant_ptq
  - numerical_verify
  - perf_benchmark
  - deploy_canary

numerical_verify:
  stage: numerical_verify
  image: $CI_REGISTRY/vqa-ci:cuda118-trt86
  script:
    - python tools/verify_numerics.py --fp32 model.onnx --int8 model_int8.onnx --trt model_int8.trt --calib_data /data/calib_500
  rules:
    - if: $CI_PIPELINE_SOURCE == "merge_request_event"
    - if: $CI_COMMIT_BRANCH == "main"

perf_benchmark:
  stage: perf_benchmark
  image: $CI_REGISTRY/vqa-ci:cuda118-trt86
  script:
    - trtexec --loadEngine=model_int8.trt --iterations=1000 --warmUp=200 --avgRuns=100 --dumpProfile --exportProfile=profile.json
    - python tools/parse_trt_profile.py profile.json --p99_threshold_ms=5 --throughput_threshold_fps=1500
  artifacts:
    reports:
      metrics: profile.json