TAOISM

KO | EN

gitlite — search

#python #react #ai #typescript #rust #linux #docker #nodejs #macos #java #windows #cli #awesome-list #ios

(★ 52)

TAOISM: A TEE-based Confidential Heterogeneous Deployment Framework for DNN Models

TAOISM 최신버젼 다운로드

최종 버전 다운로드 (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

<h1>TAOISM: A <ins>T</ins>EE-b<ins>a</ins>sed C<ins>o</ins>nfident<ins>i</ins>al Heterogeneou<ins>s</ins> Fra<ins>m</ins>ework for DNN Models</h1>
<div align="center">
  <a href="https://opensource.org/license/mit/">
    <img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-blue">
  </a>
  <a href="https://pytorch.org/">
    <img src="https://img.shields.io/badge/PyTorch->=v1.7.0-EE4C2C.svg?style=flat-square">
  </a>
  <a href="hhttps://opensource.org/license/mit/">
    <img src="https://img.shields.io/badge/SGX_SDK->2.6-green">
  </a>
</div>

TAOISM is a <ins>T</ins>EE-b<ins>A</ins>sed c<ins>O</ins>nfident<ins>I</ins>al heterogeneou<ins>S</ins> fra<ins>M</ins>ework for DNN Models. TAOISM is light-weight, flexible and compatible to PyTorch models. TAOISM can put some privacy-related and critical layers of a DNN model into SGX enclaves, and the rest of the model is executed on GPU. TAOISM is designed to evaluate the inference speed of various TEE-Shielded DNN Partition (TSDP) strategies.

<center>
    <img src="asset/TSDP.png" width=100% >
</center>
<p align="center">
    <em>Demonstration of six TSDP strategies.</em>
    <br>
</p>

For more detail of the framework and TSDP strategies, please refer to our S&P'24 paper:

[*No Privacy Left Outside: On the (In-)Security of TEE-Shielded DNN Partition for On-Device ML*](https://arxiv.org/abs/2310.07152) </br>
**Ziqi Zhang, Chen Gong, Yifeng Cai, Yuanyuan Yuan, Bingyan Liu, Ding Li, Yao Guo, Xiangqun Chen** </br>

If you found this code helpful for your research, please cite our paper:

```
@article{zhang2023no,
  title={No Privacy Left Outside: On the (In-) Security of TEE-Shielded DNN Partition for On-Device ML},
  author={Zhang, Ziqi and Gong, Chen and Cai, Yifeng and Yuan, Yuanyuan and Liu, Bingyan and Li, Ding and Guo, Yao and Chen, Xiangqun},
  journal={arXiv preprint arXiv:2310.07152},
  year={2023}
}

@inproceedings{zhang2022teeslice,
  title={TEESlice: Slicing DNN Models for Secure and Efficient Deployment},
  author={Zhang, Ziqi and Ng, Lucien KL and Liu, Bingyan and Cai, Yifeng and Li, Ding and Guo, Yao and Chen, Xiangqun},
  booktitle={Proceedings of the 2nd ACM International Workshop on AI and Software Testing/Analysis},
  pages={1--8},
  year={2022}
}
```

- [1. How to Install](#1-how-to-install)
- [2. How to run](#2-how-to-run)
- [4. Customize Your Model](#4-customize-your-model)
- [5. Code Structure](#5-code-structure)
- [6. Disclaimer](#6-disclaimer)
- [7. Acknowledgement](#7-acknowledgement)

##  1. <a name='HowtoInstall'></a>How to Install

1. Install Intel's [linux-sgx](https://github.com/intel/linux-sgx) and [linux-sgx-driver](https://github.com/intel/linux-sgx-driver).
    We tested our code on SGX SDK 2.15 To test if your hardware supports SGX, you can use this [repo](https://github.com/ayeks/SGX-hardware). To check if SGX SDK is correctly installed, you can check if `/dev/isgx` exists by running `ls /dev/isgx`. You can also run the sample code in SGX SDK to test if SGX is correctly installed.

1. Install PyTorch with Python3. You may install it using [Anaconda](https://www.anaconda.com/)

conda create -n taoism python=3.7
        conda install pytorch==1.7.0 torchvision==0.8.0 torchaudio==0.7.0 cudatoolkit=11.0 -c pytorch

2. Make the C++ part of this repo

make

NOTE: Currently the compilation's `SGX_MODE` is set to `HW`, which means the code can only run on a machine with SGX hardware. If you want to run the code on a machine without SGX hardware, you can compile in the simulation mode by changing the `SGX_MODE` to `SIM` in the Makefile.

##  2. <a name='Howtorun'></a>How to run
1. Source your Intel SGK SDK environment. For example

source /opt/intel/sgxsdk/environment

2. We provide two scripts (in `teeslice/scripts`) as a demonstration of our framework, you should run the scripts from the **root** directory of this project. (Do not enter the `teeslice` nor `teeslice/scripts` directories.)

- `teeslice/scripts/run_resnet_baseline.sh` can run a ResNet-18 model in three modes (GPU mode, CPU mode, and Enclave mode)
- `teeslice/scripts/run_teeslice.sh` runs three models trained from TEESlice on CIFAR100 dataset.

##  4. <a name='CustomizeYourModel'></a>Customize Your Model

We implement basic layers for CNN models. You can use these layers to build your own model and customize how to deployment your model. You can put some layers into SGX and the rest of the model to GPU. You can refer to `teeslice/sgx_resnet_cifar.py` and `teeslice/sgx_resnet_cifar_teeslice.py` for an example.

For each layer, you can use `EnclaveMode` parameter at initialization to specify how to deploy this layer. 
| EnclaveMode | Description |
| --- | --- |
| EnclaveMode.GPU | Run the layer on GPU |
| EnclaveMode.ENCLAVE | Run the layer in SGX |
| EnclaveMode.AUTO | Run the layer in SGX if the input tensor is in SGX, otherwise run on GPU |

NOTE: If you want to construct a network of other architectures, you may need to change the memory chunk size in the enclave. Due to the memory limitation of SGX, currently the memory in the SGX is managed in a chunk-based manner. The chunk size is defined in `/Include/common_with_enclaves.h` by `STORE_CHUNK_ELEM` and `WORK_CHUNK_ELEM` (Currently they are the same). For each layer, the framework loads one chunk in the memory and perform the operation on the chunk.

For ResNet, if your input image size is 224x224, the chunk size should be set to 9408. If your input is 32x32, the chunk size should be set to 4704. The criterion to set the chunk size is that it should compatible with the tensor size for each layer during matrix multiplication. For linear layer, the chunk size should be divisible by input channels. For convolution layer, because the current implementation is based on im2col, the chunk size should be divisible $input\_width * input\_channel$ and $output\_channel$ (as the assert statements in `SGXDNN/layers/conv.cpp`).

If you want to add new layer (take `linear` as an example) implementation in SGX, you should modify following files:
- Add the file of python object in `python/layers/sgx_xxx.py`
- Add the python interface for the initialization, forward, and backward (if you need backward) functions in `python/enclave_interfaces.py`
- Add the argument type for the python interface in `__init__` function of the object `EnclaveInterface` in `python/enclave_interfaces.py`
- Add C++ application interface in `App/enclave_bridge.cpp`. The C++ application interface calls the ECALL functions
- Add ECALL definition in `Enclave/Enclave.edl`
- Add ECALL implementation in `Enclave/sgxdnn.cpp`. The implementation calls the enclave functions
- Add the declaration and implementation of enclave functions in `*.hpp` and `*.cpp` files in `SGXDNN/layers/`

##  5. <a name='CodeStructure'></a>Code Structure

We summarize the code structure and the functionality of major files/directories to help you understand and reuse the code.

```

##  6. <a name='Disclaimer'></a>Disclaimer
DO NOT USE THIS SOFTWARE TO SECURE ANY REAL-WORLD DATA OR COMPUTATION!

This software is a proof-of-concept meant for performance testing of the TAOISM framework ONLY. It is full of security vulnerabilities that facilitate testing, debugging and performance measurements. In any real-world deployment, these vulnerabilities can be easily exploited to leak all user inputs.

Some parts that have a negligble impact on performance but that are required for a real-world deployment are not currently implemented (e.g., setting up a secure communication channel with a remote client and producing verifiable attestations).

##  7. <a name='Acknowledgement'></a>Acknowledgement
The framework is based on the code of Goten (Lucien K. L. Ng, 2021) and Slalom (Tramer & Boneh, 2019).

TAOISM

TAOISM 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

TAOISM

TAOISM 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

Download README