import torch
from torch.onnx import export

model.eval()
dummy = torch.randn(1, 3, 256, 144)
export(
    model, dummy, "portrait_seg.onnx",
    opset_version=17,
    input_names=["image"], output_names=["mask"],
    dynamic_axes={"image": {0: "batch", 2: "h", 3: "w"},
                  "mask": {0: "batch", 2: "h", 3: "w"}},
    do_constant_folding=True
)

polygraphy run portrait_seg.onnx --trt --onnxrt --compare-outputs

from tensorrt import IInt8EntropyCalibrator2, ICudaEngine

class SegCalibrator(IInt8EntropyCalibrator2):
    def __init__(self, data_loader, cache_file="calib.cache"):
        self.loader = data_loader
        self.cache_file = cache_file
        self.stream = iter(data_loader)

    def get_batch_size(self): return self.loader.batch_size
    def get_batch(self, names):
        try: return [next(self.stream).astype(np.float32).ctypes.data]
        except StopIteration: return None
    def read_calibration_cache(self):
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f: return f.read()
    def write_calibration_cache(self, cache):
        with open(self.cache_file, "wb") as f: f.write(cache)

builder = trt.Builder(logger)
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = SegCalibrator(calib_loader)
engine = builder.build_serialized_network(network, config)

config->setFlag(BuilderFlag::kFP16);           // 混合精度兜底
config->setFlag(BuilderFlag::kPREFER_PRECISION_CONSTRAINTS); // 强制精度约束
config->setPreviewFeature(PreviewFeature::kDISABLE_EXTERNAL_TACTIC_SOURCES_FOR_CORE_0805, true);
config->setTacticSources(1U << TacticSource::kCUBLAS_LT); // 仅用 cuBLASLt，减少自动调优耗时

IOptimizationProfile* profile = builder->createOptimizationProfile();
profile->setDimensions("image", OptProfileSelector::kMIN, Dims4{1,3,128,72});
profile->setDimensions("image", OptProfileSelector::kOPT, Dims4{1,3,256,144});
profile->setDimensions("image", OptProfileSelector::kMAX, Dims4{1,3,512,288});
config->addOptimizationProfile(profile);

trtexec --onnx=portrait_seg.onnx 
        --saveEngine=portrait_seg_int8.trt 
        --int8 --calib=calib.cache 
        --optShapes=image:1x3x256x144 
        --minShapes=image:1x3x128x72 
        --maxShapes=image:1x3x512x288 
        --workspace=512 --verbose 2>&1 | tee build.log

// 1. 加载参数与模型
ncnn::Net net;
net.opt.use_vulkan_compute = true;          // 优先 Vulkan
net.opt.use_fp16_packed = true;             // FP16 存储
net.opt.use_int8_inference = true;          // 开启 INT8 路径
net.load_param("portrait_seg_int8.param");
net.load_model("portrait_seg_int8.bin");

// 2. 预处理 + 推理
ncnn::Mat in = ncnn::Mat::from_pixels_resize(
    rgb_data, ncnn::Mat::PIXEL_RGB, w, h, 256, 144);
in.substract_mean_normalize(mean_vals, norm_vals);

ncnn::Extractor ex = net.create_extractor();
ex.set_num_threads(4);
ex.input("image", in);
ncnn::Mat out;
ex.extract("mask", out);  // 1 x 1 x 256 x 144

// 3. 后处理：双线性上采样至原图 + 阈值 0.5 二值化
cv::Mat mask(cv::Size(w, h), CV_32FC1, out.data);
cv::resize(mask, mask, cv::Size(w, h), 0, 0, cv::INTER_LINEAR);
cv::threshold(mask, mask, 0.5, 1.0, cv::THRESH_BINARY);

import torch.ao.quantization.quantize_fx as quantize_fx
from torch.ao.quantization import QConfigMapping, get_default_qat_qconfig

# 1. 准备 QConfig：权重 Per-Channel 对称，激活 Per-Tensor 非对称（移动端友好）
qconfig = get_default_qat_qconfig('qnnpack')  # 移动端后端参考配置
qconfig_mapping = QConfigMapping().set_global(qconfig)

# 2. 模型融合：Conv+BN+ReLU -> Fused Module（必须在 prepare_qat 前完成）
model_fused = torch.ao.quantization.fuse_modules_fx(model, [["conv", "bn", "act"]])

# 3. 插入 FakeQuant 观察器
model_prepared = quantize_fx.prepare_qat_fx(model_fused, qconfig_mapping, example_inputs)

# 4. 微调训练（建议冻结 BN 统计量，仅解冻 FakeQuant 参数）
optimizer = torch.optim.AdamW([
    {"params": [p for n, p in model_prepared.named_parameters() if "fake_quant" in n], "lr": 1e-4},
    {"params": [p for n, p in model_prepared.named_parameters() if "fake_quant" not in n], "lr": 1e-5}
])
# 仅跑 5-10 epochs，配合 Knowledge Distillation (KD) 从 FP32 Teacher 学习
for epoch in range(10):
    train_one_epoch(model_prepared, loader, optimizer, criterion_kd=nn.KLDivLoss())

# 5. 转换为部署用 INT8 图
model_int8 = quantize_fx.convert_fx(model_prepared.eval())
torch.jit.save(torch.jit.trace(model_int8, dummy), "portrait_seg_qat_int8.pt")

// 核心 Kernel 伪代码 (CUDA/HIP 统一编写，便于 Android Vulkan 移植)
__global__ void FusedPostProcessKernel(
    const half* __restrict__ logits,   // [1, 1, H/16, W/16] INT8 dequant 后 FP16
    half* __restrict__ output_mask,    // [1, 1, H, W] FP16/UINT8
    const half* __restrict__ guide_img,// [1, 3, H, W] 原图 RGB FP16 (零拷贝共享)
    int src_h, int src_w, int dst_h, int dst_w,
    float threshold, float eps, int radius
) {
    // 1. 双线性插值上采样 + Sigmoid 融合 (避免中间全局内存写入)
    // 2. 引导滤波快速实现 (O(1) 盒式滤波 + 共享内存缓存导向图)
    // 3. 阈值二值化直接写 UINT8 输出
}

import tensorrt as trt
from tensorrt.tensorrt import PluginFieldCollection, PluginField

def replace_postprocess_with_plugin(network: trt.INetworkDefinition):
    # 1. 找到最后的 Unresize/Sigmoid/Threshold 节点
    last_layer = network.get_output(0).producer
    # 2. 移除原有后处理子图
    network.unmark_output(last_layer.get_output(0))
    # 3. 插入自定义 Plugin
    plugin_creator = trt.get_plugin_registry().get_plugin_creator("FusedPostProcess", "1", "")
    pfc = PluginFieldCollection([
        PluginField("threshold", np.array([0.5], dtype=np.float32)),
        PluginField("guide_radius", np.array([4], dtype=np.int32)),
        PluginField("eps", np.array([1e-3], dtype=np.float32))
    ])
    plugin_layer = network.add_plugin_v2([last_layer.get_output(0), guide_input], plugin_creator.create_plugin("post", pfc))
    network.mark_output(plugin_layer.get_output(0))

class UnifiedBufferPool {
    struct BufferHandle {
        void* cpu_ptr;          // 宿主映射地址
        VkBuffer vk_buf;        // Vulkan Buffer
        VkDeviceMemory vk_mem;  // 绑定的 DMA-BUF FD
        size_t size;
        uint64_t frame_id;      // 双/三缓冲同步标识
    };
    std::deque<BufferHandle> free_list_, in_flight_;
    std::mutex mtx_;

public:
    // 从 Camera HAL 申请/导入 AHardwareBuffer
    BufferHandle acquire(int w, int h, int format, uint64_t frame_id) {
        std::lock_guard lk(mtx_);
        if (free_list_.empty()) return allocate_new(w, h, format);
        auto hdl = free_list_.front(); free_list_.pop_front();
        hdl.frame_id = frame_id;
        in_flight_.push_back(hdl);
        return hdl;
    }

    // 推理完成后归还，由渲染线程显式 release
    void release(BufferHandle hdl) {
        std::lock_guard lk(mtx_);
        in_flight_.erase(std::find(in_flight_.begin(), in_flight_.end(), hdl));
        free_list_.push_back(hdl);
    }
};

// 简易 Native 侧解密加载流程
ICudaEngine* load_encrypted_engine(const char* enc_path, const uint8_t* key, size_t key_len) {
    std::vector<uint8_t> enc_data = read_file(enc_path);
    std::vector<uint8_t> plain_data(enc_data.size() - 28); // GCM tag 16B + IV 12B
    aes_gcm_decrypt(key, enc_data.data()+12, enc_data.size()-28, enc_data.data(), plain_data.data());
    // 验证 SHA256
    if (sha256(plain_data) != EXPECTED_HASH) throw std::runtime_error("Integrity check failed");
    return runtime->deserializeCudaEngine(plain_data.data(), plain_data.size());
}

# 1. ONNX 导出时固定 Image Encoder 分辨率 (1024x1024)，Prompt Encoder 动态
torch.onnx.export(
    sam.image_encoder, dummy_img, "encoder.onnx",
    dynamic_axes={"image": {2: "H", 3: "W"}},  # 仅 H/W 动态，Batch=1 固定
    opset_version=17
)

# 2. TensorRT 构建时使用 Strongly Typed API 显式指定 Attention Mask 处理
builder_config.set_flag(trt.BuilderFlag.STRONGLY_TYPED)  # 强制 FP16/INT8 类型推导
profile = builder.create_optimization_profile()
profile.set_shape("image", (1,3,512,512), (1,3,1024,1024), (1,3,1536,1536))
# Prompt Encoder 稀疏输入 (points/boxes) 单独 Profile
profile.set_shape("point_coords", (1,5,2), (1,10,2), (1,20,2))

# 3. 启用 FlashAttention-2 Plugin (需 TensorRT OSS 编译)
# 注意：移动端 GPU 共享内存带宽受限，FlashAttention 仅在 L>1024 时收益为正
# 建议：Encoder 端用 Linear Attention 变体，仅 Decoder 端保留标准 MHSA

# 自动搜索每层最优精度 (INT8/FP16/FP32) 满足精度约束
autoquant --onnx=encoder.onnx 
          --calib_data=calib_set/ 
          --accuracy_metric=mIoU 
          --target_accuracy_drop=0.3% 
          --output=encoder_mixed.trt

# GitLab CI / Argo Rollouts 片段
steps:
  - setWeight: 5      # 5% 流量
  - analysis:
      templates:
      - templateName: seg-canary-analysis
        args:
        - name: new-engine-hash
          value: "{{NEW_ENGINE_HASH}}"
  - setWeight: 20
  - analysis: ...
  - setWeight: 100

# AnalysisTemplate 核心逻辑 (PromQL)
- name: latency_regression
  successCondition: |
    histogram_quantile(0.99, rate(seg_inference_latency_bucket{version="canary"}[5m])) 
    <= 
    1.1 * histogram_quantile(0.99, rate(seg_inference_latency_bucket{version="stable"}[5m]))
- name: accuracy_drop
  successCondition: |
    avg(seg_mask_iou_sampled{version="canary"}) >= 0.99 * avg(seg_mask_iou_sampled{version="stable"})

[Mobile End]                          [Cloud Edge]
  ┌─────────────┐                       ┌─────────────┐
  │ Light Encoder │── Low-res Feat (1/32) ──▶│ Heavy Refiner │
  │ (MobileViT)   │◀── High-res Delta ──────│ (Swin-L +    │
  │ INT8 2ms      │   (Diffusion/Residual)  │  Mask2Former)│
  └─────────────┘                           └─────────────┘
        │                                           │
        ▼                                           ▼
  Local Mask (Coarse)                     Cloud Mask (Fine)
        └────────────── Fusion (Guided Filter) ──────────────┘
                           │
                           ▼
                    Final 4K Mask @ 30fps