ConFlow
中文 | English
ConFlow是一个支持迁移TEE安全计算到NNA（Neural Network Accelerators，神经网络加速器）的神经网络计算框架：
-
安全计算迁移：可自动将线性算子以密文形式迁移到各类神经网络加速器上，例如GPU、NPU、TPU等，在保证数据以密文形式计算的同时，利用各类NNA的算力。
-
前端模型结构编译：自动迁移的密文算子以线性拆分的形式确保安全性，通过前端模型描述语言的编译器，转化为框架支持的密态数据结构。
- [2024.02] 我们更新了相关的文档及说明。
- [2023.12] 我们开源了模型描述编译器和适配TensorFlow的框架实现，能够将传统计算的模型描述，编译为支持此密文计算的模型结构。
- [2023.07] 我们提出了一种隐私保护机器学习框架，支持数据密态计算的同时只需增加少量的计算开销。
ConFlow 是一个支持迁移TEE安全计算到NNA（Neural Network Accelerators，神经网络加速器）的神经网络计算框架，可自动将线性算子以密文形式迁移到各类神经网络加速器上，例如GPU、NPU、TPU等，在保证数据以密文形式计算的同时，利用各类NNA的算力提高性能；同时，为了减少开发者对此类密文网络结构的适配难度，该框架自带一个模型描述语言的编译器，能将TensorFlow的模型描述自动编译到此框架的密态模型结构上。
-
🏆 隐私保护：线性算子是以密文形式在NNA上进行计算，算力提供方将无法访问到明文；而非线性算子运行在TEE内，算力提供方也无法访问到明文。整个流程可以保护隐私数据被泄露。
-
🔥 性能表现：相比较普通的TEE框架，本项目可以支持NNA的算力加速，实际吞吐量会比单独的CPU TEE执行明显要高。
在多个模型和框架上的详细评测结果。
我们将 ConFlow 应用在 Imagenet 数据集上，并存储模型中间的特征图。
需要预先安装好 TensorFlow (version > 2.0) 与 Intel SGX 相关软硬件
-
en_ops :
- Makefile
- *.cc : TensorFlow自定义算子的C++文件。
-
sgx_tf_ops :
- App ：加密算子与外界的接口
- Enclave：加密算子的真正实现文件
- Makefile
-
model_demo :
- DeepCrossing：DeepCrossing模型，将输入进行加密，并将非线性算子替换为加密算子。
-
operators.py : 加密算子的python包装层，模型可直接使用import operators进行使用。
-
run.sh : 配置及编译脚本
- 克隆我们的仓库并跳转到相应目录
git clone git@github.com:Nightliver/privacy_tf.git
cd privacy_tf- 配置密钥及编译算子
sh run.sh #运行脚本
请输入扩大倍数: #输入一个数字，该数字为数据加密后扩大的倍数
请输入抽取次数: #输入一个数字，该数字为数据加密时随机抽取的次数，例如5，为从随机生成的数据中抽取5次。
请输入密钥：#输入一串数字，用‘,’分隔开，数字数目应与抽取次数相同
#等待编译完成即可。en_ops文件夹下会出现ops.so,sgx_tf_ops文件夹下会出现sgx.so文件。测试脚本
import tensorflow.compat.v1 as tf #导入tensorflow包
import operators as sgx #导入加密算子包
sgx.create_enclave() #进行可信执行环境的初始化
a = tf.constant([[2,3,4],[-1,-2,0]],tf.float32) #创建一个tensor
b = sgx.en_crypt(a) #对tensor进行加密处理
c = sgx.e_relu(b) #对加密的tensor进行加密的relu运算
d = sgx.de_crypt(c) #对加密计算的结果进行解密cd ./model_demo/Deepcrossing #进入模型目录
python train.py #运行加密模型利用TEE框架自定义SGX算子的示例
-
在./sgx_tf_ops/Enclave/Enclave.cpp中实现具体计算逻辑：
void ecall_test(float *input,int N,float *output){ } void ecall_test_grad(float *input,float *grad,int N,float *output){ }
-
在./sgx_tf_ops/Enclave/Enclave.edl中的trusted 中注册：
public void ecall_test([user_check]float* input,int N,[user_check]float* output); public void ecall_test_grad([user_check]float* input,[user_check]float* grad,int N,[user_check]float* output); -
在./sgx_tf_ops/App/App.h中写接口的头文件：
void test(unsigned long int eid,float *input, int N, float *output); void test_grad(unsigned long int eid,float *input,float* grad, int N, float *output);
-
在./sgx_tf_ops/App/App.cpp中写接口函数：
void test(unsigned long int eid,float *input, int N, float *output){ sgx_status_t ret = ecall_test(eid, input,N,output); if (ret != SGX_SUCCESS) { print_error_message(ret); throw ret; } } void test_grad(unsigned long int eid,float *input,float *grad, int N, float *output){ sgx_status_t ret = ecall_test_grad(eid, input,grad,N,output); if (ret != SGX_SUCCESS) { print_error_message(ret); throw ret; } }
-
最后重新make编译即可。
构建自定义TensorFlow算子的示例
-
在./en_ops文件夹下新建test.cc文件
-
导入头文件以及设置命名空间：
#include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/shape_inference.h" #include <cstdlib> #include <iostream> #include <cmath> #include <dlfcn.h> using namespace tensorflow; using namespace shape_inference; using namespace std;
-
注册ops
REGISTER_OP("ETest") .Input("input: float") .Attr("eid_low: int") .Attr("eid_high: int") .Attr("times:int") .Output("output: float") .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { c->set_output(0, c->input(0)); //形状设定 return Status::OK(); }); REGISTER_OP("ETestGrad") .Input("grad: float") .Input("input: float") .Attr("eid_low: int") .Attr("eid_high: int") .Attr("times:int") .Output("output: float");
-
注册kernels
class ETestOp : public OpKernel { public: explicit ETestOp(OpKernelConstruction* context) : OpKernel(context) { OP_REQUIRES_OK(context, context->GetAttr("eid_low", &eid_low_)); OP_REQUIRES_OK(context, context->GetAttr("eid_high", &eid_high_)); OP_REQUIRES_OK(context, context->GetAttr("times", ×_)); lib = dlopen("(所在路径)/sgx.so",RTLD_LAZY); OP_REQUIRES(context, lib != NULL, errors::Unknown("Unable to load sgx.so!")); } void Compute(OpKernelContext* context) override { const Tensor& input = context->input(0); //获取输入 auto input_flat = input.flat<float>(); const TensorShape& input_shape = input.shape(); //获取输入形状 Tensor* output = NULL; OP_REQUIRES_OK(context, context->allocate_output(0, input_shape, &output)); //初始化输出 auto output_flat = output->flat<float>(); int N = input_flat.size()/times_; unsigned long int eid_ = (eid_high_ << 32) + eid_low_; typedef void (*function)(unsigned long int eid,float* input, int N, float* output); dlerror(); function test_kernel = (function) dlsym(lib, "test"); //调用SGX算子 const char *dlsym_error = dlerror(); OP_REQUIRES(context, !dlsym_error, errors::Unknown("loading of test failed: ", dlsym_error)); //失败处理 test_kernel(eid_,(float*)input_flat.data(),N,(float*)output_flat.data());//调用SGX进行计算 }; private: void* lib; int64 eid_low_; int64 eid_high_; int64 times_; }; REGISTER_KERNEL_BUILDER(Name("ETest").Device(DEVICE_CPU), ETestOp); class ETestGradOp : public OpKernel { public: explicit ETestGradOp(OpKernelConstruction* context) : OpKernel(context) { OP_REQUIRES_OK(context, context->GetAttr("eid_low", &eid_low_)); OP_REQUIRES_OK(context, context->GetAttr("eid_high", &eid_high_)); OP_REQUIRES_OK(context, context->GetAttr("times", ×_)); lib = dlopen("(所在路径)/sgx.so",RTLD_LAZY); OP_REQUIRES(context, lib != NULL, errors::Unknown("Unable to load sgx.so!")); } void Compute(OpKernelContext* context) override { const Tensor& grad = context->input(0); const Tensor& input = context->input(1); auto grad_flat = grad.flat<float>(); auto input_flat = input.flat<float>(); // check shapes of input const TensorShape& input_shape = input.shape(); // create output tensor Tensor* output = NULL; OP_REQUIRES_OK(context, context->allocate_output(0, input_shape, &output)); auto output_flat = output->flat<float>(); const int N = input_flat.size()/times_; unsigned long int eid_ = (eid_high_ << 32) + eid_low_; typedef void (*function)(unsigned long int eid,float* input,float *grad, int N, float* output); dlerror(); function test_grad_kernel = (function) dlsym(lib, "test_grad"); const char *dlsym_error = dlerror(); OP_REQUIRES(context, !dlsym_error, errors::Unknown("loading of relu_grad failed: ", dlsym_error)); test_grad_kernel(eid_,(float*)input_flat.data(),(float*)grad_flat.data(),N,(float*)output_flat.data()); }; private: void* lib; int64 eid_low_; int64 eid_high_; int64 times_; }; REGISTER_KERNEL_BUILDER(Name("ETestGrad").Device(DEVICE_CPU), ETestGradOp);
-
使用make重新编译
-
python层包装及梯度注册
-
在operators.py中进行函数包装和梯度注册：
def e_test(inputs):
global eid,times
return trainer_ops.e_test(inputs,eid_low=(eid&0xFFFFFFFF),eid_high=(eid>>32),times=times)
@ops.RegisterGradient("ETest")
def _e_test_grad(op, grad):
with tf.name_scope("ETestGrad"), tf.xla.experimental.jit_scope(compile_ops=False):
return trainer_ops.e_test_grad(grad,op.inputs[0],eid_low=op.get_attr("eid_low"), eid_high=op.get_attr("eid_high"),times = op.get_attr("times"))使用自定义的加密算子的示例
import tensorflow as tf
import operators
operators.create_enclave() #初始化enclave
operators.e_test() #使用自定义算子- 本仓库中代码依照 Apache-2.0 协议开源，并遵照本项目相关声明。
- 直接或间接使用本项目原创的相关方法得到学术与商业等利益，项目原作者拥有知情权。
如果您觉得我们模型/代码/论文有帮助，请给我们 ⭐ 和 引用 📝，感谢！
@INPROCEEDINGS{10247964,
author={Li, Qiushi and Ren, Ju and Zhang, Yan and Song, Chengru and Liao, Yiqiao and Zhang, Yaoxue},
booktitle={2023 60th ACM/IEEE Design Automation Conference (DAC)},
title={Privacy-Preserving DNN Training with Prefetched Meta-Keys on Heterogeneous Neural Network Accelerators},
year={2023},
volume={},
number={},
pages={1-6},
keywords={Training;Privacy;Data privacy;Design automation;Prefetching;Artificial neural networks;Resists;deep learning;privacy preserving;neural network accelerator;cloud computing;trusted execution environment},
doi={10.1109/DAC56929.2023.10247964}}