python 对 xxx.pt或者xxx.pth文件加密和解密

在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')  

原文地址：https://blog.csdn.net/shanglianlm/article/details/140602376

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！