如何防止模型被窃取？基于TVM的端上模型部署加密方案

TVM简介

TVM全称Tensor Virtual Machine，中文名称张量虚拟机，它主要的用途是将深度学习模型编译成最小可部署的模块，并提供更多的后端自动生成和优化模型的基础设施。架构如下：

它的特点是：

1.  最佳的性能：通过编译和最小运行时，在现有的硬件上优化机器学习工作负载，进而发挥最佳性能。

2.  丰富的后端：可以在CPU、GPU、浏览器、微控制器、FPGA等硬件上运行，同时可自动在多种后端生成和优化张量算子。

3.  灵活可用：TVM 的灵活设计支持对区块稀疏度、量化（1、2、4、8 位整数，posit）、随机森林/经典 ML、内存规划、MISRA-C 兼容性、Python 原型设计。

4.  简单易用：在 Keras、MXNet、PyTorch、Tensorflow、CoreML、DarkNet 等的深度学习模型上编译。用 Python 借助 TVM 进行编译，即可直接使用 C++、Rust 或 Java 完成推理部署。

TVM工作流程

TVM属于领域编译器，其工作流程也遵循通用CPU编译器的设计规则，将复杂的任务简单化，并进行多层抽象， TVM也是采用多层IR设计，IR的转换如下：

多层IR转换以后，最后生成指定的后端文件，如上图所示TVM编译会生成三个文件：mod.json是模型的计算图描述文件，mod.so是经过优化生成的算子，mod.param是存储的模型参数文件，其文件格式如下：

Safe Tensor

我们提出了Safe Tensor的概念。Safe Tensor的设计原则是简单高效易用，我们放弃了全量加密、mpc等技术，使用了传统的对称加解密算法，在TVM生成mod.param文件的时候，在其保存每层tensor的过程中将里面的数据进行加密。与原始mod.param文件比较，我们只加密了最关键的数据，在加解密速度上保证了最快，也将引入的加解密将损耗降到了最低。Safe Tensor采用了传统的对称加密算法。为了确保业务安全，密钥的保存和使用应该遵循最佳实践。在使用密钥之前，需要使用非对称加密算法进行解密。 

TVM生成模型参数文件格式就变成如下：

效果如下：

Safe Tensor的方案只是加密了模型最重要的tensor数据，这对TVM的上层是无感的，运行的方法和之前介绍的方法也相同。

Safe Tensor的方案有如下优势：

1.  相对TVM生成的推理程序，加密后业务接入无感，灵活易用。

2.  推理时性能无损失。

3.  模型文件被加密，和推理代码一一绑定，直接拷走模型文件后不可用，从而起到保护作用。

Safe Tensor给大家提供了一个新的模型保护思路，现阶段只是在模型文件生成和加载的时候进行加解密。如果想获得更好的保护效果，还需要进程内存保护的技术保证内存dump相关的攻击。随着未来AI技术的发展，硬件设备厂商也会慢慢关注这块市场，也会推出相关的硬件特性，届时可以使用软硬结合的安全保护，让模型应用既高效又安全。