如何利用计算图切分实现可信环境（TEE）与非安全环境（REE）协同推理

在端侧AI应用中，保护模型权重或用户隐私数据至关重要。传统的全加密推理（如全同态加密）性能极差，而“计算图切分”技术提供了一种实用的平衡方案：将涉及敏感隐私的计算环节（如首层特征提取或特定包含关键权重的中间层）放置在硬件隔离的可信执行环境（TEE）中，而将计算密集型且不敏感的层留在非安全环境（REE）中使用GPU/NPU加速。本文将详解如何使用 PyTorch FX 进行自动化的计算图切分。

1. 技术背景

• REE (Rich Execution Environment)：如 Android/Linux 系统，拥有强大的算力资源，但安全性较低。
• TEE (Trusted Execution Environment)：如 ARM TrustZone 或国产芯片的安全核，提供硬件级隔离，安全性极高但算力极度受限。

协同调度的核心在于：在图编译阶段识别出“安全边界”，并插入通信存根，实现跨环境的数据流转。

2. 核心实操：基于 PyTorch FX 的图切分

PyTorch FX 是一个强大的图分析与变换工具，我们可以利用它将模型逻辑切分为多个 Submodule。

步骤一：定义包含隐私敏感层的模型

import torch
import torch.nn as nn
from torch.fx import symbolic_trace

class PrivacyModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 假设第一层涉及敏感参数（例如人脸识别的原始模板提取）
        self.sensitive_layer = nn.Conv2d(3, 16, 3, padding=1)
        self.relu = nn.ReLU()
        # 后续层为计算密集型，可放在 REE
        self.compute_heavy_layer = nn.Sequential(
            nn.Conv2d(16, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(32, 10)
        )

    def forward(self, x):
        x = self.sensitive_layer(x)
        x = self.relu(x)
        x = self.compute_heavy_layer(x)
        return x

步骤二：识别边界并自动化切分

我们将利用 FX 追踪计算图，并将其拆分为 TEE 部分和 REE 部分。

def split_model_for_tee(model):
    traced = symbolic_trace(model)
    # 生产环境中可根据 node.op 或 node.name 搜索切分点
    # 这里演示手动逻辑：切分出敏感层及其后的激活层

    class TEEPart(nn.Module):
        def __init__(self, original):
            super().__init__()
            self.layer = original.sensitive_layer
            self.relu = original.relu
        def forward(self, x): 
            return self.relu(self.layer(x))

    class REEPart(nn.Module):
        def __init__(self, original):
            super().__init__()
            self.heavy = original.compute_heavy_layer
        def forward(self, x): 
            return self.heavy(x)

    return TEEPart(model), REEPart(model)

步骤三：模拟跨环境协同调度

在推理时，由于 TEE 和 REE 内存不互通，需要显式的内存拷贝过程。

def collaborative_inference(input_data, tee_module, ree_module):
    # 1. 在 TEE 环境中执行（受保护，不可见）
    with torch.no_grad():
        # 模拟 TEE 内部操作
        intermediate_feature = tee_module(input_data)

    # 2. 跨环境通信桥接
    # 实际底层会调用类似 TEE_CopyOutSharedMemory 的 API
    # 这里使用 detach 模拟从受控内存导出的过程
    ree_input = intermediate_feature.detach().clone()

    # 3. 在 REE 环境执行计算密集型层（GPU 加速）
    with torch.no_grad():
        final_output = ree_module(ree_input)

    return final_output

# 运行测试
model = PrivacyModel()
tee, ree = split_model_for_tee(model)
data = torch.randn(1, 3, 32, 32)
result = collaborative_inference(data, tee, ree)
print(f\"Inference finished. Output shape: {result.shape}\")

3. 关键优化点

1. 最小化切分面：跨 TEE/REE 切换存在上下文切换（Context Switch）开销。切分点应选在张量尺寸较小的瓶颈处（如 Pooling 层后），以减少内存带宽压力。
2. 安全性加固：TEE 导出的张量建议进行简单的混淆或去相关性操作，防止攻击者通过 REE 中的特征图进行模型逆向攻击。
3. 异构流水线：利用多缓冲区技术（Double Buffering），在 REE 处理当前帧 N 的密集计算时，TEE 可以提前开始处理下一帧 N+1 的敏感层，从而隐藏切换时延。

4. 总结

计算图切分是端侧 AI 安全落地的必经之路。通过将逻辑结构转化为物理隔离的子图，开发者可以在不牺牲大部分性能的前提下，有效保障模型资产和用户数据的安全。