欢迎光临散文网 会员登陆 & 注册

TensorRT 模型加密杂谈

2023-08-09 19:38 作者:laugh12321  | 我要投稿

在大多数项目交付场景中,经常需要对部署模型进行加密。模型加密一方面可以防止泄密,一方面可以便于模型跟踪管理,防止混淆。

由于博主使用的部署模型多为TensorRT格式,这里以TensorRT模型为例,讲解如何对模型进行加密、解密以及推理加密模型。

代码仓库:https://github.com/laugh12321/TRTCrypto

加密算法的选择和支持的库

Crypto++ 是C/C++的加密算法库,基本上涵盖了市面上的各类加密解密算法,包括对称加密算法(AES等)和非对称加密算法(RSA等)。

两种算法使用的场景不同,非对称加密算法一般应用于数字签名和密钥协商的场景下,而对称加密算法一般应用于纯数据加密场景,性能更优。在对模型的加密过程中使用对称加密算法。

简易版本

以AES-GCM加密模式为例,编写一个检测的加密、解密方法

上述代码,使用AES-CBC加密模式对数据进行加密、解密,其中key和iv为加密算法的参数,keySize为key的长度。

加密流程:

  1. 初始化加密器,设置key和iv

  2. 读取文件内容并存储在字符串data中

  3. 使用加密器对data进行加密,加密后的内容存储在字符串cipher中

解密流程:

  1. 初始化解密器,设置key和iv

  2. 读取加密后的文件内容并存储在字符串cipher中

  3. 使用解密器对cipher进行解密,解密后的内容存储在字符串recovered中

转换为序列化格式

推理加密模型的方法有两种,一种是将模型解密后保存为文件再进行推理,另一种是将模型解密后转换为序列化格式,再进行推理。 很明显第一种方式比较鸡肋,因为每次推理都需要进行解密,而且解密后的模型文件也会暴露在外面,不安全。这里使用第二种方式,将模型解密后转换为序列化格式进行推理。这里给出一个简单的例子,将存储解密数据的字符串recovered进行序列化。

上述代码的返回值可以直接作为TensorRT的推理引擎的输入。例如进行反序列化操作

使用MAC地址作为密钥

一般情况下,我们只想让客户在指定的机器上运行模型,这时候就需要使用机器的唯一标识作为密钥,这里使用MAC地址作为密钥。

上述代码,使用了不同的API获取MAC地址,其中Windows使用GetAdaptersInfo函数,Linux使用getifaddrs函数。

有了MAC地址,就可以将MAC地址转换为密钥,这里使用SHA256算法对MAC地址进行哈希,然后取前32个字节作为密钥。

添加头部信息

为了新的文件能够被区分和可迭代,除了加密后的数据外还添加了头部信息,比如为了判断该文件类型使用固定的魔数作为文件的开头;为了便于后面需求迭代写入版本号以示区别;为了能够在解密时判断是否采用了相同的密钥将加密时的密钥进行SHA256计算后存储;这三部分构成了目前加密后文件的头部信息。加密后的文件包含头部信息 + 密文信息。

上述代码中,加密函数EncryptWithHeader中,首先计算密钥的SHA256哈希值,然后将魔数、版本号、密钥哈希值、密文依次拼接,然后使用AES-256-CBC算法加密,最后将头部信息和密文拼接返回。解密函数DecryptWithHeader中,首先从密文中提取出头部信息,然后使用密钥哈希值验证密钥是否正确,最后使用AES-256-CBC算法解密密文。

参考资料

  • PaddleX-模型加密部署

  • MegEngine-模型加解密和打包

  • 《机器学习系统:设计和实现》-模型的安全保护

  • Crypto++ GCM Mode


TensorRT 模型加密杂谈的评论 (共 条)

分享到微博请遵守国家法律