在Python中,直接对.pt或.pth(PyTorch模型文件)进行加密和解密通常不是直接由PyTorch库本身支持的,因为这些文件主要是用于存储模型的权重和参数,而不是设计为加密存储。然而,你可以通过读取这些文件的内容,使用Python的加密库(如cryptography)来加密这些内容,然后将加密后的数据存储到新的文件中。解密过程则相反,从加密的文件中读取内容,解密后再保存到原始格式的文件中。

以下是一个简单的示例,展示了如何使用cryptography库来加密和解密.pt文件(注意,这里假设.pt文件的内容可以被当作二进制数据进行处理):

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes  
from cryptography.hazmat.primitives import padding  
from cryptography.hazmat.backends import default_backend  
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC  
from os import urandom  
import base64  
from torch import load, save  
  
def generate_key(password: str, salt: bytes, key_length: int = 32) -> bytes:  
    kdf = PBKDF2HMAC(  
        algorithm=hashes.SHA256(),  
        length=key_length,  
        salt=salt,  
        iterations=100000,  
        backend=default_backend()  
    )  
    key = kdf.derive(password.encode())  
    return key  
  
def encrypt_file(input_file: str, output_file: str, password: str):  
    # 读取原始文件  
    with open(input_file, 'rb') as file_in:  
        data = file_in.read()  
  
    # 生成密钥和IV  
    salt = urandom(16)  
    key = generate_key(password, salt)  
    iv = urandom(16)  # 初始化向量  
  
    # 加密数据  
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv))  
    encryptor = cipher.encryptor()  
    padder = padding.PKCS7(128).padder()  
    padded_data = padder.update(data) + padder.finalize()  
    encrypted_data = encryptor.update(padded_data) + encryptor.finalize()  
  
    # 将salt和IV与加密数据一起写入输出文件  
    with open(output_file, 'wb') as file_out:  
        file_out.write(salt)  
        file_out.write(iv)  
        file_out.write(encrypted_data)  
  
def decrypt_file(input_file: str, output_file: str, password: str):  
    # 读取加密文件  
    with open(input_file, 'rb') as file_in:  
        salt = file_in.read(16)  
        iv = file_in.read(16)  
        encrypted_data = file_in.read()  
  
    # 生成密钥  
    key = generate_key(password, salt)  
  
    # 解密数据  
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv))  
    decryptor = cipher.decryptor()  
    padded_data = decryptor.update(encrypted_data) + decryptor.finalize()  
    unpadder = padding.PKCS7(128).unpadder()  
    data = unpadder.update(padded_data) + unpadder.finalize()  
  
    # 写入解密后的数据到文件  
    with open(output_file, 'wb') as file_out:  
        file_out.write(data)  
  
# 假设你有一个PyTorch模型文件叫做'model.pt'  
encrypt_file('model.pt', 'model_encrypted.pt', 'your_secret_password')  
# 解密后,你可以将'model_decrypted.pt'重新加载为PyTorch模型  
decrypt_file('model_encrypted.pt', 'model_decrypted.pt', 'your_secret_password')  
Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐