DarkneTZ

项目概述

DarkneTZ 是一个基于 OP-TEE 的神经网络安全推理框架，实现了模型分割和内存保护。支持分类器训练和推理，将指定层加载入 TEE (可信执行环境) 以保护模型隐私和安全性。

代码修复

由于 DarkneTZ 项目较老，存在一些代码语法问题需要处理。以下列出所需的修改：

1. classifier.c 文件修改

文件路径：optee_examples/darknetz/host/examples/classifier.c

修改内容：在第 192 行和 777 行修改为：

2. darknet.c 文件修改

文件路径：optee_examples/darknetz/host/examples/darknet.c

修改内容：

• 在第 5 行下增加一行：#include <darknet.h>
• 在第 408 行左右进行修改，将 return 0 改为 return

3. convolutional_layer.c 文件修改

文件路径：optee_examples/darknetz/host/src/convolutional_layer.c

修改内容：在第 475 行进行修改：

4. darknet.h 文件修改

文件路径：optee_examples/darknetz/host/include/darknet.h

修改内容：在第 7 行增加如下声明：

5. mnist_lenet.cfg 文件修改

文件路径：out-br/target/root/cfg/mnist_lenet.cfg

修改内容：更新网络配置文件如下：

编译步骤

完成上述修改后，进行编译步骤如下：

该命令将进入 OP-TEE 构建目录，并使用 16 个并行线程进行编译和运行。

核心特性

• TEE 集成：基于 OP-TEE 框架实现可信执行环境
• 模型分割：支持指定层数范围保存在 TEE 中
• 安全推理：保护模型参数不被恶意访问
• MNIST 示例：完整的手写数字识别实现
• 前后向传播：TEE 与 REE 间的完整训练流程

技术栈

• OP-TEE - 开放式可信执行环境
• C 语言 - 核心实现
• 神经网络 - 深度学习框架
• MNIST - 数据集标准

环境测试

OP-TEE 环境测试（xtest）：

DarkneTZ 项目分类器训练（MNIST 手写数字数据集）

推理测试

DarkneTZ 项目学习

以分类器（classifier）为例进行详细分析。

模型架构

cfg/mnist_lenet.cfg 为 MNIST 卷积神经网络模型架构，含若干节（section），由 [ ] 标识，每节有若干配置项（option），由键值对表示。

• 第一节必为 [net] 或 [network]，包含神经网络的总体配置，如每批次训练数量(batch)，学习率(learning_rate)等。
• 之后各节为各层神经网络配置，如 [convolutional] 卷积层，[maxpool] 最大池化层，[connected] 全连接层，[dropout] 随机丢弃层，[softmax] 层等，各层的节下配置了输出维度（output），卷积核大小(size)，激活函数(activation)等具体架构参数。

数据集

cfg/mnist.dataset 含有分类器（classifier）的训练数据默认配置，如 10 类（classes）即 0 到 9，训练数据列表路径(train)，标签列表路径（labels）和名称列表路径（names），topK（top）等。

分类器的训练数据集为 MNIST 手写数字训练数据集，图像数据存储在 data/mnist/images 目录下，命名格式为 <type>_<index>_<label>.png，如 t_00002_c4.png 表示 train（训练）数据集的第 00002 张图片，图像为数字 4。

推测 valid 为交叉验证(cross validation)数据集，test 为测试数据集，但相应目录下可能存在数据缺失情况。

模型训练

通过命令进行模型训练：

使用数据集 mnist.dataset 训练 mnist_lenet.cfg 架构的卷积神经网络分类器模型，其中第 4 层至第 10 层存入 TEE 以保护。

train_classifier

• 首先通过配置文件加载网络，分配空间，根据配置进行参数的初始化，并将分割点间的层加载入 TEE。
• 随后进行加载数据，通过数据加载线程每次并发加载 net->batch * net->subdivisions * ngpus（每批次数量*分支*gpu数量）张图片，存入缓冲区 buffer（train）。

train_network

通过加载的数据训练神经网络。TEE 与 CEE 通过 net_truth_CA 通信，通过 train_network_datum 进行前向传播->反向传播->更新网络参数，然后返回损失。

forward_network（前向传播）

根据各层类型选择相应函数在各层进行前向传播：

• Dropout 层：置零输入或随机缩放以防止过拟合。
• Connected 层：进行矩阵计算或归一化。
• Convolutional 层：将图像转为列向量，然后分组卷积矩阵运算。
• Maxpool 层：提取区域内最大值实现降维。
• Cost 层：位于最后，根据损失函数类型计算与真实值的损失，示例为 sse（均方误差）。

各层推理结果 output 通过神经网络 net 的 input 传递给下一层，softmax 层为实际输出层，计算图像识别各数字概率。

对于 TEE 中的层，通过函数 forward_network_CA 以在 TEE 中前向传播，然后通过 forward_network_back_CA 接收 TEE 中神经网络的推理结果。

Cost 层的输出可用于反向传播调整参数。最后计算成本，calc_network_cost 将各层成本平均值存入 net.cost。

backward_network（反向传播）

根据各层类型，激活函数类型，在各层进行反向传播，梯度计算。

在 TEE 与 REE 分界处，使用：

• backward_network_back_CA_addidion
• backward_network_back_CA
• backward_network_CA

函数向 TA 传入输入，在 TEE 执行反向传播；使用 backward_network_CA_addidion 函数从 TA 传回结果，在 REE 继续反向传播。

update_network（更新参数）

需最后再同时更新所有参数而非过程中更新（梯度下降法需要）。

模型推理

通过命令进行推理：

加载训练的权重 mnist_lenet.weights 进行推理，其中 4 至 10 层于 TEE 中。

在推理时的输出 top 可能出现乱码：

参考 GitHub Issue #37 修改 predict_classifier 函数的 free 位置。

predict_classifier

加载网络配置和权重，解析参数，加载并处理需要推理的图片，随后调用 network_predict 进行预测，输出 top 个结果。

network_predict

进行前向传播 forward_network，并获取结果，当最后层在 TEE 时通过 net_output_return_CA 函数获取结果。

TA 与 CA 的通信

除一般性的会话和上下文的初始化和销毁之外，TA 与 CA 间的通信体现在网络加载、前向传播与反向传播过程中，在 DarkneTZ 项目中 TA 与 CA 交互相关函数名称含有标识，在 CA 的部分标识"CA"，在 TA 的部分标识"TA"。

同一函数在 TA 与 CA 内部的不同版本执行相似操作，分隔点由 OP-TEE 命令和共享内存传参，安全性由 OP-TEE 系统机制保证。

网络加载

• Load_network -> parse_network_cfg -> make_network 在 CA 中加载网络
• Parse_network_cfg -> parse_net_options -> make_network_CA 在 TA 中加载分隔点间的网络

传播过程

• 在分隔点从 CA 传播入 TA 时，构建参数内存空间并使用 OP-TEE 命令传入 TA
• 在分隔点从 TA 传播至 CA 时，构建共享内存空间并使用 OP-TEE 命令从 TA 传出

相关链接

GitHub 仓库 · 相关Issue