YOLOv8模型加密与解密部署

YOLOv8模型加密与解密部署

前言

最近项目需要，需要对训练好的模型加密，并且完成解密部署，趁着yolov8的发布写篇博客，我这里主要用python实现，C++版的也在开发中，有机会也分享出来。

模型加密的手段有很多种，目前现有的加解密工具也有很多，我目前使用的以openssl、cryptography为主。这些工具里都提供了很多加密算法如AES等等。

对于训练好的模型而言，加解密的主要流程分为：生成密钥（密钥的格式可以是公司名+模型+加版本号+发布日期，也可以随机生成）、由密钥将例如pt格式的模型加密并保存为加密文件、最后读取加密文件和密钥并解密出模型。

解密出模型后将模型数据存放到内存中（保存到某个变量里），然后由推理框架提供的从内存中加载模型的接口来加载模型，这样才是一个比较完备的加解密流程。

YOLOv8的加解密

yolov8是ultralytics公司的开源项目，由于yolov8的代码耦合度比较高，直接在源码的基础上修改会涉及到很多部分，很容易出bug且不一定能修改成功，因此我这里选择把pt模型转为onnx，然后用ORT实现推理的流程，并融入加解密的部分。

1.加密示例

加密部分的示例代码如下：给定onnx文件和要保存的加密文件以及密钥即可生成加密后的模型。

def model_encryption(pth_file, encryp_file, license):
    with open(pth_file, 'rb') as fr:
        pth_bytes = fr.read()

    key = read_license(license)
    cipher_suite = Fernet(key)
    encrypted_data = cipher_suite.encrypt(pth_bytes)

    with open(encryp_file, 'wb') as fw:
        fw.write(encrypted_data)

保存结果如下图，pt是官网的预训练模型，yolov8n-en是加密后的文件

2.解密示例

解密部分的示例代码如下：读取解密文件并将解密之后数据流保存到变量

def model_decryption(encryt_file, license):
    with open(encryt_file, 'rb') as fr:
        encrypted_data = fr.read()
    key = read_license(license)
    decrypted_data = Fernet(key).decrypt(encrypted_data)
    return decrypted_data

后续ORT直接加载模型推理即可。

推理结果如下，模型的精度基本没有丢失。

欢迎交流入群

在CV行业有一定经验的学者、对CV感兴趣的朋友、或是CV行业的辛苦打工人欢迎加入交流群，刚建的新群名额少快快加入，群友会给你及时的分享CV技术！
群名额少快快加入，群友会给你及时的分享CV技术！
（群号：937412537）