# export_onnx.py
import torch
from model import PortraitSegNet  # 自定义轻量化骨干网络

def export_onnx(weight_path, onnx_path, input_shape=(1, 3, 256, 256)):
    model = PortraitSegNet(num_classes=2)
    model.load_state_dict(torch.load(weight_path, map_location='cpu'))
    model.eval()

    dummy = torch.randn(*input_shape)
    torch.onnx.export(
        model, dummy, onnx_path,
        opset_version=13,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}},
        do_constant_folding=True
    )
    print(f"ONNX exported to {onnx_path}")

if __name__ == '__main__':
    export_onnx('best.pth', 'portrait_seg.onnx')

# calib_dataset.py
import cv2, numpy as np
from torch.utils.data import Dataset

class CalibDataset(Dataset):
    def __init__(self, img_dir, size=(256, 256), max_samples=500):
        self.paths = [f"{img_dir}/{f}" for f in os.listdir(img_dir)][:max_samples]
        self.size = size

    def __len__(self): return len(self.paths)

    def __getitem__(self, idx):
        img = cv2.imread(self.paths[idx])
        img = cv2.resize(img, self.size)
        img = img.astype(np.float32) / 255.0
        img = (img - 0.5) / 0.5  # 归一化与训练一致
        return img.transpose(2, 0, 1)  # CHW

trtexec --onnx=portrait_seg_sim.onnx 
        --calib=calib_dataset.py 
        --int8 
        --calibTable=portrait_seg_int8.cache 
        --minShapes=input:1x3x256x256 
        --optShapes=input:1x3x256x256 
        --maxShapes=input:4x3x256x256 
        --workspace=1024 
        --saveEngine=portrait_seg_int8.trt

// TrtBuilder.cpp
#include <NvInfer.h>
#include <fstream>

using namespace nvinfer1;

ICudaEngine* buildEngine(const std::string& onnxPath, const std::string& cachePath,
                         bool useFP16, bool useINT8, size_t workspaceGB = 1) {
    IBuilder* builder = createInferBuilder(gLogger);
    IBuilderConfig* config = builder->createBuilderConfig();
    config->setMaxWorkspaceSize(workspaceGB << 30);

    if (useFP16) config->setFlag(BuilderFlag::kFP16);
    if (useINT8) {
        config->setFlag(BuilderFlag::kINT8);
        IInt8Calibrator* calibrator = new EntropyCalibrator(cachePath);
        config->setInt8Calibrator(calibrator);
    }

    IParser* parser = createParser(*network, gLogger);
    parser->parseFromFile(onnxPath.c_str(), static_cast<int>(ILogger::Severity::kWARNING));

    // 显式批次维度优化
    IOptimizationProfile* profile = builder->createOptimizationProfile();
    profile->setDimensions("input", OptProfileSelector::kMIN, Dims4{1,3,256,256});
    profile->setDimensions("input", OptProfileSelector::kOPT, Dims4{1,3,256,256});
    profile->setDimensions("input", OptProfileSelector::kMAX, Dims4{4,3,256,256});
    config->addOptimizationProfile(profile);

    ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    // 序列化保存
    IHostMemory* serialized = engine->serialize();
    std::ofstream ofs("portrait_seg.trt", std::ios::binary);
    ofs.write(static_cast<char*>(serialized->data()), serialized->size());
    return engine;
}

// TrtInfer.cpp
ICudaEngine* loadEngine(const std::string& enginePath, IRuntime* runtime) {
    std::ifstream ifs(enginePath, std::ios::binary);
    ifs.seekg(0, std::ios::end);
    size_t size = ifs.tellg();
    ifs.seekg(0, std::ios::beg);
    std::vector<char> buffer(size);
    ifs.read(buffer.data(), size);
    return runtime->deserializeCudaEngine(buffer.data(), size, nullptr);
}

// 推理单帧
void infer(IExecutionContext* ctx, float* d_input, float* d_output, int batch) {
    void* bindings[] = {d_input, d_output};
    ctx->enqueueV2(bindings, stream, nullptr);
    cudaStreamSynchronize(stream);
}

# qat_prepare.py
import torch
from torch.ao.quantization import QConfigMapping, get_default_qat_qconfig
from torch.ao.quantization.quantize_fx import prepare_qat_fx

def prepare_qat_model(model, example_inputs):
    # 1. 仅对 Conv/Linear/Act 插入 FakeQuantize，跳过 BN（已融合）、Resize、Concat
    qconfig = get_default_qat_qconfig('qnnpack')  # 移动端后端对齐
    qconfig_mapping = QConfigMapping().set_global(qconfig)
    
    # 2. 自定义例外：分割头最后一层 1x1 Conv 保持 FP32，保护类别边界精度
    class_names = [n for n, _ in model.named_modules()]
    for name in class_names:
        if 'seg_head.final_conv' in name:
            qconfig_mapping.set_module_name(name, None)  # 禁用量化
    
    model_prepared = prepare_qat_fx(model, qconfig_mapping, example_inputs)
    return model_prepared

# 导出 QAT 模型（保留 FakeQuant 节点）
torch.onnx.export(model_qat, dummy, 'portrait_qat.onnx', opset_version=17, 
                  operator_export_type=torch.onnx.OperatorExportTypes.ONNX_FALLTHROUGH)

# trtexec 直接读取 ONNX 中的 scale/zero-point，无需再跑校准
trtexec --onnx=portrait_qat.onnx --int8 --saveEngine=portrait_qat_int8.trt

// PortraitPreprocessPlugin.cpp
#include <cuda_fp16.h>
#include <NvInferPlugin.h>

class PortraitPreprocessPlugin : public nvinfer1::IPluginV3, 
                                 public nvinfer1::IPluginV3OneCore,
                                 public nvinfer1::IPluginV3OneBuild,
                                 public nvinfer1::IPluginV3OneRuntime {
public:
    // 构造函数接收 mean/std，序列化存入
    PortraitPreprocessPlugin(const float* mean, const float* std) { ... }
    PortraitPreprocessPlugin(const void* data, size_t length) { ... } // 反序列化

    // 核心计算：单 Kernel 融合 HWC->CHW + Normalize + Pad
    int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc,
                    nvinfer1::PluginTensorDesc const* outputDesc,
                    void const* const* inputs, void* const* outputs,
                    void* workspace, cudaStream_t stream) noexcept override {
        launch_preprocess_kernel(
            static_cast<const uint8_t*>(inputs[0]),  // NHWC u8
            static_cast<float*>(outputs[0]),         // NCHW fp32/fp16
            inputDesc[0].dims.d[1], inputDesc[0].dims.d[2], // H, W
            mMean, mStd, mTargetH, mTargetW, stream
        );
        return 0;
    }

    // 形状推导：支持动态输入分辨率
    nvinfer1::DimsExprs getOutputDimensions(int32_t outputIndex,
        const nvinfer1::DimsExprs* inputs, int32_t nbInputs,
        nvinfer1::IExprBuilder& exprBuilder) noexcept override {
        return nvinfer1::DimsExprs{4, {inputs[0].d[0], exprBuilder.constant(3), 
                                       exprBuilder.constant(mTargetH), exprBuilder.constant(mTargetW)}};
    }

    // 必须实现：clone, getNbOutputs, getOutputDataType, supportsFormatCombination...
    // 注册宏
    REGISTER_TENSORRT_PLUGIN(PortraitPreprocessPluginCreator);
};

# onnx_register_plugin.py
import torch.onnx.symbolic_helper as sym_help
from torch.onnx import register_custom_op_symbolic

@parse_args('v', 'v', 'v', 'v', 'v', 'i', 'i')
def symbolic_portrait_preprocess(g, input, mean, std, target_h, target_w, interp_mode, align_corners):
    return g.op("com.company::PortraitPreprocess", input, 
                mean_f=mean, std_f=std, target_h_i=target_h, target_w_i=target_w)

register_custom_op_symbolic('torchvision::normalize', symbolic_portrait_preprocess, 13)

Camera (YUV_420_888) 
  → ImageReader (AHardwareBuffer) 
  → libyuv NV21→RGBA (GPU Shader / Renderscript) 
  → EGLImageKHR → cudaGraphicsEGLRegisterImage 
  → CUDA Kernel (Preprocess Plugin 直读纹理) 
  → TensorRT 推理 
  → CUDA Kernel (Postprocess 直写纹理) 
  → EGLImage → SurfaceView / TextureView 渲染

// ZeroCopyManager.cpp
cudaGraphicsResource_t registerEGLImage(EGLDisplay dpy, EGLImageKHR img) {
    cudaGraphicsResource_t res;
    cudaGraphicsEGLRegisterImage(&res, dpy, img, cudaGraphicsRegisterFlagsSurfaceLoadStore);
    return res;
}

void mapAndInfer(cudaGraphicsResource_t res, IExecutionContext* ctx, cudaStream_t stream) {
    cudaArray_t arr; cudaGraphicsMapResources(1, &res, stream);
    cudaGraphicsSubResourceGetMappedArray(&arr, res, 0, 0);
    // 绑定到 TensorRT 输入 Binding
    void* d_ptr; cudaGraphicsResourceGetMappedPointer(&d_ptr, nullptr, res);
    ctx->setTensorAddress("input", d_ptr);
    ctx->enqueueV3(stream);
    cudaGraphicsUnmapResources(1, &res, stream);
}

// PipelineScheduler.cpp
struct FrameContext { cudaStream_t stream; cudaEvent_t done; void* bindings[2]; };

void runPipeline(std::vector<FrameContext>& frames, int numFrames) {
    for (int i = 0; i < numFrames; ++i) {
        int cur = i % 2, next = (i+1) % 2;
        // 1. 等待上一帧 H2D 完成
        cudaStreamWaitEvent(frames[cur].stream, frames[next].done, 0);
        // 2. 启动推理
        ctx->enqueueV3(frames[cur].stream);
        cudaEventRecord(frames[cur].done, frames[cur].stream);
        // 3. 并行在 Host 提交下一帧 H2D (或使用 cudaGraph 进一步降低启动开销)
        submitNextFrameH2D(frames[next].stream); 
    }
}

// ContextPool.h
class ContextPool {
    std::vector<std::unique_ptr<IExecutionContext>> pool_;
    std::mutex mtx_; std::condition_variable cv_;
public:
    ContextPool(ICudaEngine* engine, int maxContexts) {
        for (int i=0; i<maxContexts; ++i) 
            pool_.push_back(std::unique_ptr<IExecutionContext>(engine->createExecutionContext()));
    }
    std::unique_ptr<IExecutionContext, std::function<void(IExecutionContext*)>> acquire() {
        std::unique_lock lk(mtx_);
        cv_.wait(lk, [this]{ return !pool_.empty(); });
        auto ctx = std::move(pool_.back()); pool_.pop_back();
        return {ctx.release(), [this](IExecutionContext* c){ release(c); }};
    }
    void release(IExecutionContext* ctx) {
        std::lock_guard lk(mtx_); pool_.push_back(std::unique_ptr<IExecutionContext>(ctx));
        cv_.notify_one();
    }
};

// ModelCrypto.cpp
#include <openssl/evp.h>

std::vector<uint8_t> encryptEngine(const std::vector<uint8_t>& plain, const std::string& deviceId) {
    // Key = HKDF(master_key, deviceId)  // 一机一密
    // IV = random 12 bytes
    // AAD = "TensorRT_Engine_v1" | model_hash
    // 返回: IV(12) || Ciphertext || Tag(16)
}

std::vector<uint8_t> decryptAndVerify(const std::vector<uint8_t>& blob, const std::string& deviceId) {
    // 解密 + GCM Tag 验证，失败即拒绝加载
}

CDN (版本清单 manifest.json) 
  → 客户端定时轮询 (指数退避) 
  → 规则引擎匹配 (设备型号/OS/ABI/当前版本/灰度比例) 
  → 增量下载 (bsdiff 差分包，平均 2-5 MB) 
  → 校验签名 (Ed25519) + 解密 
  → 原子替换 (rename + fsync) 
  → 下次启动生效 / 支持热重载 (dlclose/dlopen 新 engine)

# daily_inspection.py
def run_inspection():
    # 1. 从特征仓库拉取 200 张代表性难例 (逆光、遮挡、复杂背景)
    # 2. 在标杆机型 (3 高/中/低端) 跑推理
    # 3. 对比基线引擎：延迟、显存、输出余弦相似度
    # 4. 生成 HTML 报告 + 发送企业微信/钉钉卡片
    # 5. 若延迟回归 > 5% 或 相似度 < 0.995 → 阻断发布流水线

# 1. 详细性能剖析 (生成 JSON 供 Perfetto/Chrome Tracing 打开)
trtexec --engine=portrait_int8.trt --profilingVerbosity=detailed --exportProfile=profile.json

# 2. 逐层数值对比 (需准备 FP32/INT8 两个 engine 及同一输入 dump)
trtexec --engine=portrait_fp32.trt --dumpOutput --output=out_fp32.bin
trtexec --engine=portrait_int8.trt --dumpOutput --output=out_int8.bin
python compare_tensors.py out_fp32.bin out_int8.bin --tol=1e-3

# 3. 显存占用快照
nvidia-smi dmon -s pucvmet -d 1 -o TD

# 4. CUDA Graph 捕获验证
nsys profile --trace=cuda,nvtx --cuda-graph-trace=node ./trtexec --engine=... --graph=portrait.graph

# 5. 模型结构可视化 (Netron 支持 .trt 直接拖拽)
# 或转 ONNX 再看: trtexec --engine=... --exportOutput=model.onnx