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)
YOLOv8模型加密与解密部署_第1张图片

你可能感兴趣的:(pytorch,深度学习,python深度学习,YOLO,python,深度学习)