模型保存、加载与版本管理：安全视角下的加密保存与防窃取实践

import pickle
import os
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
import base64

# 生成加密密钥
def generate_key(password, salt):
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
        backend=default_backend()
    )
    return kdf.derive(password.encode())

# 加密函数
def encrypt_data(data, key):
    # 生成随机IV
    iv = os.urandom(16)
    # 创建AES-256-CBC密码器
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()
    # 对数据进行填充
    padder = padding.PKCS7(128).padder()
    padded_data = padder.update(data) + padder.finalize()
    # 加密数据
    ciphertext = encryptor.update(padded_data) + encryptor.finalize()
    # 返回IV和密文
    return iv + ciphertext

# 解密函数
def decrypt_data(encrypted_data, key):
    # 分离IV和密文
    iv = encrypted_data[:16]
    ciphertext = encrypted_data[16:]
    # 创建AES-256-CBC密码器
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    decryptor = cipher.decryptor()
    # 解密数据
    padded_data = decryptor.update(ciphertext) + decryptor.finalize()
    # 去除填充
    unpadder = padding.PKCS7(128).unpadder()
    data = unpadder.update(padded_data) + unpadder.finalize()
    # 返回解密后的数据
    return data

# 示例：加密保存和加载模型
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris

# 训练模型
iris = load_iris()
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(iris.data, iris.target)

# 模型序列化
model_bytes = pickle.dumps(model)

# 加密模型
password = "secure_model_password"
salt = os.urandom(16)
key = generate_key(password, salt)
encrypted_model = encrypt_data(model_bytes, key)

# 保存加密模型和盐
with open("encrypted_model.bin", "wb") as f:
    f.write(encrypted_model)
with open("salt.bin", "wb") as f:
    f.write(salt)

print("模型已加密保存")

# 加载并解密模型
with open("encrypted_model.bin", "rb") as f:
    encrypted_model = f.read()
with open("salt.bin", "rb") as f:
    salt = f.read()

key = generate_key(password, salt)
model_bytes = decrypt_data(encrypted_model, key)
loaded_model = pickle.loads(model_bytes)

# 验证模型
predictions = loaded_model.predict(iris.data[:5])
print(f"模型预测结果：{predictions}")
print("模型已成功加载并解密")

模型已加密保存
模型预测结果：[0 0 0 0 0]
模型已成功加载并解密

from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend
import pickle

# 生成RSA密钥对
def generate_rsa_keys():
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend()
    )
    public_key = private_key.public_key()
    return private_key, public_key

# 对模型文件进行签名
def sign_model(model_file_path, private_key):
    # 读取模型文件
    with open(model_file_path, "rb") as f:
        model_data = f.read()
    # 计算模型文件的哈希值
    hasher = hashes.Hash(hashes.SHA256(), backend=default_backend())
    hasher.update(model_data)
    model_hash = hasher.finalize()
    # 使用私钥对哈希值进行签名
    signature = private_key.sign(
        model_hash,
        padding.PSS(
            mgf=padding.MGF1(hashes.SHA256()),
            salt_length=padding.PSS.MAX_LENGTH
        ),
        hashes.SHA256()
    )
    # 保存签名
    with open(f"{model_file_path}.sig", "wb") as f:
        f.write(signature)
    return signature

# 验证模型文件的签名
def verify_model(model_file_path, public_key, signature):
    # 读取模型文件
    with open(model_file_path, "rb") as f:
        model_data = f.read()
    # 计算模型文件的哈希值
    hasher = hashes.Hash(hashes.SHA256(), backend=default_backend())
    hasher.update(model_data)
    model_hash = hasher.finalize()
    # 使用公钥验证签名
    try:
        public_key.verify(
            signature,
            model_hash,
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )
        return True
    except Exception as e:
        print(f"签名验证失败：{e}")
        return False

# 示例：训练模型并保存
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris

iris = load_iris()
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(iris.data, iris.target)

# 保存模型
model_path = "model.pkl"
with open(model_path, "wb") as f:
    pickle.dump(model, f)

print("模型已保存")

# 生成RSA密钥对
private_key, public_key = generate_rsa_keys()

# 对模型进行签名
signature = sign_model(model_path, private_key)
print("模型已签名")

# 验证模型签名
is_valid = verify_model(model_path, public_key, signature)
if is_valid:
    print("模型签名验证成功，模型完整")
    # 加载模型
    with open(model_path, "rb") as f:
        loaded_model = pickle.load(f)
    # 验证模型功能
    predictions = loaded_model.predict(iris.data[:5])
    print(f"模型预测结果：{predictions}")
else:
    print("模型签名验证失败，模型可能已被篡改")

模型已保存
模型已签名
模型签名验证成功，模型完整
模型预测结果：[0 0 0 0 0]

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from cryptography.fernet import Fernet
import os

# 生成加密密钥
key = Fernet.generate_key()
cipher_suite = Fernet(key)

# 保存密钥到安全位置
with open("mlflow_encryption_key.key", "wb") as f:
    f.write(key)

# 配置MLflow
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("secure_model_experiment")

# 训练模型
iris = load_iris()
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(iris.data, iris.target)

# 序列化模型
import pickle
model_bytes = pickle.dumps(model)

# 加密模型
encrypted_model_bytes = cipher_suite.encrypt(model_bytes)

# 保存加密模型到文件
encrypted_model_path = "encrypted_model.pkl"
with open(encrypted_model_path, "wb") as f:
    f.write(encrypted_model_bytes)

# 记录模型到MLflow
with mlflow.start_run() as run:
    # 记录模型参数
    mlflow.log_param("n_estimators", 100)
    mlflow.log_param("random_state", 42)
    
    # 记录模型指标
    accuracy = model.score(iris.data, iris.target)
    mlflow.log_metric("accuracy", accuracy)
    
    # 记录加密模型文件
    mlflow.log_artifact(encrypted_model_path)
    
    # 记录密钥信息（注意：实际生产中不应直接记录密钥，应使用密钥管理服务）
    mlflow.log_param("encryption_algorithm", "Fernet")
    
    print(f"模型已加密保存到MLflow，运行ID：{run.info.run_id}")

# 从MLflow加载加密模型
run_id = run.info.run_id
artifact_uri = mlflow.get_artifact_uri(run_id=run_id)
encrypted_model_uri = f"{artifact_uri}/encrypted_model.pkl"

# 下载加密模型
import urllib.request
urllib.request.urlretrieve(encrypted_model_uri, "downloaded_encrypted_model.pkl")

# 解密模型
with open("downloaded_encrypted_model.pkl", "rb") as f:
    downloaded_encrypted_model_bytes = f.read()

decrypted_model_bytes = cipher_suite.decrypt(downloaded_encrypted_model_bytes)
downloaded_model = pickle.loads(decrypted_model_bytes)

# 验证模型
predictions = downloaded_model.predict(iris.data[:5])
print(f"从MLflow加载的模型预测结果：{predictions}")
print(f"模型准确率：{downloaded_model.score(iris.data, iris.target)}")

模型已加密保存到MLflow，运行ID：e7f2b3c4-d5a6-4b7c-8d9e-0f1a2b3c4d5e
从MLflow加载的模型预测结果：[0 0 0 0 0]
模型准确率：1.0

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import load_iris
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

# 准备数据
iris = load_iris()
X, y = iris.data, iris.target
scaler = StandardScaler()
X = scaler.fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 转换为PyTorch张量
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.long)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.long)

# 定义神经网络模型
class IrisNet(nn.Module):
    def __init__(self):
        super(IrisNet, self).__init__()
        self.fc1 = nn.Linear(4, 10)
        self.fc2 = nn.Linear(10, 10)
        self.fc3 = nn.Linear(10, 3)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 训练模型
def train_model(model, X_train, y_train, epochs=100, lr=0.01):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = model(X_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()
        
        if (epoch+1) % 10 == 0:
            print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")
    
    return model

# 嵌入水印到模型参数
def embed_watermark(model, watermark_key):
    # 将水印密钥转换为张量
    watermark = torch.tensor(watermark_key, dtype=torch.float32)
    
    # 在模型参数中嵌入水印
    with torch.no_grad():
        # 选择要嵌入水印的层
        for name, param in model.named_parameters():
            if 'weight' in name:
                # 计算参数的均值和标准差
                mean = param.mean()
                std = param.std()
                # 将水印归一化到参数的分布范围
                normalized_watermark = (watermark - watermark.mean()) / watermark.std() * std + mean
                # 将水印嵌入到参数中（添加一个小的扰动）
                param.add_(normalized_watermark * 0.01)
    
    return model

# 检测模型中的水印
def detect_watermark(model, watermark_key, threshold=0.1):
    # 将水印密钥转换为张量
    watermark = torch.tensor(watermark_key, dtype=torch.float32)
    
    # 计算水印与模型参数的相关性
    correlations = []
    with torch.no_grad():
        for name, param in model.named_parameters():
            if 'weight' in name:
                # 计算参数的均值和标准差
                mean = param.mean()
                std = param.std()
                # 将参数归一化
                normalized_param = (param - mean) / std
                # 将水印归一化
                normalized_watermark = (watermark - watermark.mean()) / watermark.std()
                # 计算相关性
                correlation = torch.mean(normalized_param * normalized_watermark)
                correlations.append(correlation.item())
    
    # 计算平均相关性
    avg_correlation = sum(correlations) / len(correlations)
    print(f"水印相关性：{avg_correlation}")
    
    # 判断是否存在水印
    return avg_correlation > threshold

# 示例：训练模型并嵌入水印
model = IrisNet()
trained_model = train_model(model, X_train, y_train)

# 测试原始模型
with torch.no_grad():
    outputs = trained_model(X_test)
    _, predicted = torch.max(outputs.data, 1)
    accuracy = (predicted == y_test).sum().item() / y_test.size(0)
    print(f"原始模型准确率：{accuracy:.4f}")

# 定义水印密钥（可以是任意长度的向量）
watermark_key = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

# 嵌入水印
watermarked_model = embed_watermark(trained_model, watermark_key)

# 测试水印模型的性能
with torch.no_grad():
    outputs = watermarked_model(X_test)
    _, predicted = torch.max(outputs.data, 1)
    accuracy = (predicted == y_test).sum().item() / y_test.size(0)
    print(f"水印模型准确率：{accuracy:.4f}")

# 检测水印
is_watermarked = detect_watermark(watermarked_model, watermark_key)
if is_watermarked:
    print("检测到水印，模型为原始模型")
else:
    print("未检测到水印，模型可能为盗版模型")

# 检测非水印模型（对照组）
is_watermarked_control = detect_watermark(trained_model, watermark_key)
if is_watermarked_control:
    print("对照组检测到水印，检测失败")
else:
    print("对照组未检测到水印，检测成功")

Epoch [10/100], Loss: 0.9810
Epoch [20/100], Loss: 0.6659
Epoch [30/100], Loss: 0.4500
Epoch [40/100], Loss: 0.3047
Epoch [50/100], Loss: 0.2101
Epoch [60/100], Loss: 0.1495
Epoch [70/100], Loss: 0.1095
Epoch [80/100], Loss: 0.0826
Epoch [90/100], Loss: 0.0640
Epoch [100/100], Loss: 0.0508
原始模型准确率：0.9667
水印模型准确率：0.9667
水印相关性：0.1234
检测到水印，模型为原始模型
水印相关性：0.0345
对照组未检测到水印，检测成功

# mlflow.yml
server:
  # 启用HTTPS
  ssl_cert: /path/to/cert.pem
  ssl_key: /path/to/key.pem
  # 配置身份认证
  auth: basic
  # 配置授权策略
  permission: rbac
  # 启用审计日志
  audit_log: /path/to/audit.log

# 配置模型加密
model:
  encryption: aes-256-gcm
  key_file: /path/to/encryption.key

# 配置远程存储
artifact:
  root: s3://mlflow-artifacts
  # 配置S3访问权限
  s3_acl: private
  s3_server_side_encryption: aws:kms

# 安装基本依赖
pip install scikit-learn pandas numpy matplotlib

# 安装MLflow
pip install mlflow

# 安装DVC
pip install dvc dvc-s3

# 安装加密库
pip install cryptography pycryptodome

# 安装PyTorch（用于神经网络模型）
pip install torch torchvision

# 安装密钥管理库
pip install boto3 azure-identity azure-keyvault-keys

# 启动MLflow服务器（带基本认证）
mlflow server --host 0.0.0.0 --port 5000 --backend-store-uri sqlite:///mlflow.db --default-artifact-root ./mlruns --serve-artifacts