对深度学习模型权重进行加密与解密

引言

在现代深度学习应用中，模型权重（如 .pt 或 .pth 文件）是核心资产。为了保护这些模型权重防止泄露或未经授权的使用，通常需要对其进行加密和解密操作。本篇博客将介绍一种基于 AES-CBC 加密 和 HMAC 校验 的简单实现方法，确保模型文件的安全性。

1. 为什么需要对模型权重进行加密

随着深度学习的广泛应用，模型权重已经成为企业和研究机构的重要资产。未经保护的模型文件可能面临以下风险：

1. 知识产权泄露：模型权重的训练通常需要投入大量的计算资源和数据。如果权重被泄露，竞争对手可以轻松复现产品功能，造成经济和技术损失。
2. 数据篡改风险：未加密的模型文件可能在传输或存储过程中被恶意篡改，影响最终的推理结果。
3. 权限管理需求：在多用户或分布式环境中，仅授权用户应当能够访问和使用模型权重。

因此，对模型权重进行加密处理是保护知识产权、确保数据完整性的重要手段。

2. 技术简介

我们采用两种安全性技术来保护模型文件：AES-CBC 加密 和 HMAC 校验。

2.1 AES-CBC 加密

AES（Advanced Encryption Standard，高级加密标准）是一种对称加密算法，广泛应用于数据加密场景。AES 的特点是：

• 快速高效：AES 是现代加密算法中性能最优的之一。
• 高安全性：支持多种密钥长度（128/192/256 位），抗暴力破解能力强。

在 AES 加密中，有多种工作模式（如 ECB、CBC、CTR）。为了增强数据的保密性，我们选择了 CBC (Cipher Block Chaining) 模式，它的特点是：

• 引入初始向量 (IV)：在加密过程中，每个加密块会使用前一个块的密文和初始向量（IV）作为输入，这使得同样的数据每次加密的结果都不同。
• 保证加密的随机性：即使输入内容完全相同，由于 IV 的存在，加密后的输出也会不同。

AES-CBC 的加密流程：

1. 将数据分块，每块大小与 AES 的块大小（128 位，16 字节）一致。
2. 第一块数据与 IV 进行异或操作后加密，生成第一个密文块。
3. 后续每一块数据与前一个密文块进行异或后加密，依次生成剩余的密文块。

解密流程：

1. 解密每个密文块，恢复异或前的数据。
2. 使用前一个密文块或 IV 还原原始明文。

2.2 HMAC 校验

HMAC（Hash-based Message Authentication Code，基于哈希的消息认证码）是一种数据完整性校验技术，主要用于验证数据是否被篡改。HMAC 的工作原理：

1. 密钥参与计算：通过将密钥与数据结合生成一个哈希值，确保只有持有密钥的一方才能生成正确的校验码。
2. 校验完整性：在加密数据的末尾附加 HMAC 校验码，接收方可以通过重新计算 HMAC 并比对校验码，验证数据是否被篡改。

在本实现中，我们使用 SHA-256 作为哈希算法生成 HMAC 校验码。

AES-CBC 与 HMAC 的结合：

• AES-CBC 负责加密，确保机密性。
• HMAC 校验码负责验证完整性，确保数据未被篡改。

3. 加密与解密实现

接下来，我们基于 Python 使用 pycryptodome 库实现模型文件的加密和解密功能。

3.1 加密代码

以下是 AES-CBC 加密和 HMAC 校验的实现代码：

import hashlib
from pathlib import Path
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Hash import HMAC, SHA256


class ModelEncrypter:
    """模型加密与解密工具"""
    def __init__(self, key: str) -> None:
        self.key = self.hash_key(key)  # 生成 16 字节的密钥

    def encrypt(self, model_path: str) -> str:
        """加密模型并保存为 .secret 文件"""
        # 加载模型数据
        model_data = self.load_model(model_path)
        
        # 加密模型数据
        cipher = AES.new(self.key, AES.MODE_CBC)
        iv = cipher.iv  # 生成随机 IV
        encrypted_model = pad(model_data, AES.block_size)
        encrypted_model = cipher.encrypt(encrypted_model)

        # 拼接 IV 和密文
        bin_model = iv + encrypted_model

        # 生成 HMAC 校验码
        hmac = HMAC.new(self.key, digestmod=SHA256)
        hmac.update(bin_model)
        hmac_digest = hmac.digest()

        # 拼接加密数据和 HMAC
        final_data = bin_model + hmac_digest

        # 保存加密文件
        secret_path = self.get_secret_path(model_path)
        self.save_file(secret_path, final_data)
        print(f"加密模型已保存到: {secret_path}")
        return secret_path

    def decrypt(self, secret_path: str) -> bytes:
        """解密 .secret 文件并返回模型数据"""
        # 加载加密文件数据
        encrypted_data = self.load_model(secret_path)

        # 分离 IV、密文 和 HMAC
        iv = encrypted_data[:AES.block_size]
        hmac_digest = encrypted_data[-32:]
        encrypted_model = encrypted_data[AES.block_size:-32]

        # 校验 HMAC
        hmac = HMAC.new(self.key, digestmod=SHA256)
        hmac.update(encrypted_data[:-32])
        hmac.verify(hmac_digest)  # 若校验失败将抛出异常

        # 解密模型数据
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        model_data = unpad(cipher.decrypt(encrypted_model), AES.block_size)
        print("解密成功，模型数据已提取")
        return model_data

    @staticmethod
    def hash_key(key: str) -> bytes:
        """生成 16 字节的密钥"""
        sha_signature = hashlib.sha256()
        sha_signature.update(key.encode())
        return sha_signature.digest()[:16]

    @staticmethod
    def load_model(path: str) -> bytes:
        """加载文件"""
        with open(path, 'rb') as f:
            return f.read()

    @staticmethod
    def get_secret_path(path: str) -> str:
        """生成加密文件路径，后缀为 .secret"""
        return str(Path(path).with_suffix('.secret'))

    @staticmethod
    def save_file(path: str, data: bytes) -> None:
        """保存文件"""
        with open(path, 'wb') as f:
            f.write(data)

3.2 使用示例

以下是如何加密和解密模型文件的示例代码：

if __name__ == '__main__':
    # 加密密钥
    key = '159753456852'
    
    # 模型文件路径
    model_path = '/data/models/det_model.pt'
    
    # 创建加密器实例
    encrypter = ModelEncrypter(key=key)

    # 加密模型文件
    encrypted_path = encrypter.encrypt(model_path)

    # 解密模型文件
    decrypted_model = encrypter.decrypt(encrypted_path)
    
    # 可将解密后的模型数据保存为新文件（可选）
    with open('/data/models/det_model_decrypted.pt', 'wb') as f:
        f.write(decrypted_model)
    print("解密后的模型已保存")

4. 总结

通过本文的实现，我们可以：

• 安全地加密深度学习模型文件，防止泄露和滥用。
• 使用 HMAC 校验确保加密数据的完整性，防止篡改。
• 通过 AES-CBC 加密增强加密的随机性和安全性。

这套加密与解密机制可以适用于多种场景，如深度学习模型的分发、存储和授权访问。希望这篇博客对你的模型保护方案有所帮助！