Python中对AI模型的加密

在机器学习应用部署中，保护模型知识产权至关重要。本文介绍一种使用AES-256加密算法保护PyTorch模型权重文件的方法，通过强加密防止未授权访问，同时保持模型加载流程的便捷性。

直接上代码：

# encrypt_models.py
import hashlib
import io
import torch
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes

def encrypt_model(model_path: str, key: str, output_path: str):
    """
    加密模型文件并保存为加密格式
    
    参数:
        model_path: 原始模型文件路径
        key: 加密密码
        output_path: 加密后输出路径
    """
    with open(model_path, 'rb') as f:
        model_data = f.read()

    # 使用SHA-256生成256位密钥
    key_hash = hashlib.sha256(key.encode()).digest()
    
    # 生成随机初始化向量(IV)
    iv = get_random_bytes(16)
    
    # 创建AES-CBC加密器
    cipher = AES.new(key_hash, AES.MODE_CBC, iv=iv)
    
    # 加密并填充数据
    ciphertext = cipher.encrypt(pad(model_data, AES.block_size))

    # 写入IV+密文
    with open(output_path, 'wb') as f:
        f.write(iv + ciphertext)

def decrypt_model(encrypted_data: bytes, key: str) -> bytes:
    """
    解密模型数据
    
    参数:
        encrypted_data: 加密的模型数据
        key: 解密密码
        
    返回:
        解密后的原始字节数据
    """
    key_hash = hashlib.sha256(key.encode()).digest()
    
    # 分离IV和密文
    iv = encrypted_data[:16]
    ciphertext = encrypted_data[16:]
    
    # 创建AES-CBC解密器
    cipher = AES.new(key_hash, AES.MODE_CBC, iv=iv)
    
    # 解密并移除填充
    return unpad(cipher.decrypt(ciphertext), AES.block_size)

加密模型文件：

# 设置安全密钥（实际使用中应从安全来源获取）
SECRET_KEY = "Your_Strong_Password_123!"

# 加密原始模型文件
encrypt_model(
    "./weights/original_model.pth",
    SECRET_KEY,
    "./weights/encrypted_model.enc"
)

解密并加载模型：

# 读取加密模型文件
with open("./weights/encrypted_model.enc", 'rb') as f:
    encrypted_data = f.read()

# 解密模型数据
decrypted_bytes = decrypt_model(encrypted_data, SECRET_KEY)

# 创建类文件对象
file_like = io.BytesIO(decrypted_bytes)

# 安全加载模型
model = torch.load(
    file_like,
    map_location='cpu',
    weights_only=True  # 启用安全模式
)

当前密钥放在Python代码的SECRET_KEY变量中，那么实际项目中对密钥同样需要加密，我应用管理方案有以下几种：