【YOLOV5从开发到部署实战】九、YOLO训练出的模型权重加解密方法

这几年以来，AI 的概念已经火遍大江南北，能用上计算机的各个领域几乎都能用 AI 助力，对于人类来说无论是不喜的还是不用的应该都有所耳闻。那么今天提到的神经网络，就属于人工智能的分支之一机器学习的子集。具体的专业知识可以百度、知乎里找，本专栏不再赘述（up主也懒得介绍 鉴于up可怜的知识储备也无法系统性地讲解明白）。

就应用方面，比较常见的神经网络模型文件格式有 .pt/.pth (PyTorch框架)，.weights(.cfg) (Darknet框架)，.caffemodel(Caffe框架)，.pb(Tensorflow框架) ，.onnx(微软设计的一种针对机器学习所设计的开放式的文件格式)。有些朋友愿意公布自己训练的神经网络模型，那么没问题。但有些朋友做的是商业化或者单纯不想自己辛辛苦苦标注训练出的模型被别人用，这就需要加密了。

当然我也有过这方面的迷茫，相信不少人也在或者有寻找过如何加密的方式。其实模型文件就是个二进制文件 (binary file)，而加密二进制文件的方式要从网上找就轻松多了。

比较入门的级别，就是文件重命名。推理框架在读取神经网络模型时方式之一就是从文件读取，那么只要文件名路径名对了就能顺利读取。鉴于Python打包后文件众多的特点(尤其是dll文件)，完全可以取一个像模像样的文件名隐藏于众多动态链接库之中，这样别人找也会费点功夫。

更常用且安全性更高的方式，就是将模型这个二进制文件加密，需要推理时读取并解密到内存再给推理框架读取。

我们以官方YOLOv5s.pt代码为例，对模型训练的权重进行加密，加密后的结果为：

注意key.txt保存的是解密用的密钥，一定要保存好！！

加密代码

from cryptography.fernet import Fernet

# 密钥生成后自行以安全方式保存
key = Fernet.generate_key()
f = Fernet(key)

# 保存解码密钥
file = open('./key.txt','w')
file.write(str(key.decode('utf-8')))

# 未加密文件
origin_files = ['../weights/yolov5s.pt']

# 加密后的文件
encrypted_files = ['./yolov5s-en.data']

# 加密
for i in range(len(origin_files)):
    with open(encrypted_files[i],'wb') as ew:
        content = open(origin_files[i],'rb').read()
        encrypted_content = f.encrypt(content)
        ew.write(encrypted_content)

解密代码

from cryptography.fernet import Fernet

# 读取密钥
with open('./key.txt', 'rb') as key:
    key = key.readline()
f = Fernet(key)

# 加密文件
encrypted_files = ['./yolov5s-en.data']

# 解密文件
decrypt_files = ['./yolov5s-de.pt']

# 解密
for i in range(len(encrypted_files)):
    with open(decrypt_files[i],'wb') as ew:
        content = open(encrypted_files[i],'rb').read()
        decrypt_content = f.decrypt(content)
        ew.write(decrypt_content)