Python深度学习之Keras模型转换成ONNX模型流程详解

目录

• 从Keras转换成PB模型
• 从PB模型转换成ONNX模型
• 改变现有的ONNX模型精度
• 部署ONNX 模型
• 总结

从Keras转换成PB模型

请注意，如果直接使用Keras2ONNX进行模型转换大概率会出现报错，这里笔者曾经进行过不同的尝试，最后都失败了。

所以笔者的推荐的情况是：首先将Keras模型转换为TensorFlow PB模型。

那么通过tf.keras.models.load_model()这个函数将模型进行加载，前提是你有一个基于h5格式或者hdf5格式的模型文件，最后再通过改变模型的保存格式save_format参数改为tf。即可实现从Keras转换为TensorFow的格式文件了。

import tensorflow as tf
model_path = './models/model.h5'                    # 模型文件
model = tf.keras.models.load_model(model_path)
model.save('tfmodel', save_format='tf')

转换后的模型文件结构是这样的：

.
├── assets
├── keras_metadata.pb
├── saved_model.pb
└── variables
    ├── variables.data-00000-of-00001
    └── variables.index

2 directories, 4 files

从PB模型转换成ONNX模型

从PB模型转换为ONNX模型是很简单的，通过调用tf2onnx这个模块下的convert来进行处理。

你需要做的，只需要将--save-model来指定你已经转换好的TensorFlow模型，使用--output来指定你的ONNX模型输出的路径（需要指定一个独立的文件，如./xx/xx.onnx）

python -m tf2onnx.convert --saved-model ./tfmodel/ --output ./models/model.onnx --opset 11 --verbose

这样我们就得到一个ONNX模型：

.
├── model.h5
├── model.onnx
└── model_fp16.onnx

改变现有的ONNX模型精度

考虑到在不同的计算设备上，半精度和双精度锁带来的性能提升是显而易见的。

这里我使用了一个VGG16的模型来测试了fp16和fp32的性能。

----------------------
VGG Full Precision:
    Data Size: 124
    VGGFullPrecision Timing: 7.462206602096558 Seconds
    Connections: 1824812148
----------------------
VGG Half Precision:
    Data Size: 124
    VGGHalfPrecision Timing(In TensorRT): 2.563319444656372 Seconds
    Connections: 1824812148
----------------------

可以看到，在我这张RTX2060上，启用fp16相较于fp32的性能提升接近3倍。

那么我们该如何将现有的ONNX模型从fp32模型转换成fp16模型呢？

首先我们需要准备一个叫onnxmltools的库。可以通过pip来进行安装。

pip install onnxmltools

确认安装好onnxmltools后，我们通过如下的一段脚本进行精度的转换：

import onnxmltools
# 加载float16_converter转换器
from onnxmltools.utils.float16_converter import convert_float_to_float16
# 使用onnxmltools.load_model()函数来加载现有的onnx模型
# 但是请确保这个模型是一个fp32的原始模型
onnx_model = onnxmltools.load_model('../module/models/model.onnx')
# 使用convert_float_to_float16()函数将fp32模型转换成半精度fp16
onnx_model_fp16 = convert_float_to_float16(onnx_model)
# 使用onnx.utils.save_model()函数来保存，
onnxmltools.utils.save_model(onnx_model_fp16, '../module/models/model_fp16.onnx')

部署ONNX 模型

在部署ONNX模型阶段，我们将使用onnxruntime这个模块。

针对你所将使用的计算设备，如果你是CPU用户，那么你需要使用如下的指令来安装onnxruntime

pip install onnxruntime

反之，如果你的计算设备是是GPU，那么你需要使用如下的指令来安装onnxruntime

pip install onnxruntime-gpu

确认好onnxruntime安装完成后，你只需要使用如下的指令来加载你的ONNX模型即可

import onnxruntime as ort
# 指定onnx模型所在的位置
model_path = './module/models/model.onnx'
# 创建providers参数列表
providers = [
		# 指定模型可用的CUDA计算设备参数
        ('CUDAExecutionProvider', {
        	# 因为这里笔者只有一张GPU，因此GPU ID序列就为0
            'device_id': 0,
            # 这里网络额外策略使用官方默认值
            'arena_extend_strategy': 'kNextPowerOfTwo',
            # 官方这里默认建议的GPU内存迭代上限是2GB，如果你的GPU显存足够大
            # 可以将这里的2修改为其它数值
            'gpu_mem_limit': 2 * 1024 * 1024 * 1024,
            # cudnn转换算法的调用参数设置为完整搜索
            'cudnn_conv_algo_search': 'EXHAUSTIVE',
            # 确认从默认流进行CUDA流赋值
            'do_copy_in_default_stream': True,
        }),
        'CPUExecutionProvider',
    ]
# 使用onnxruntime.InferenceSession()函数创建Session
# 第一参数为模型所在的路径，第二参数为模型的providers参数列表
session = ort.InferenceSession(model_path, providers=providers)
# 通过get_input()函数和get_output()函数获取网络的输入和输出名称
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
# 使用session.run()函数执行ONNX任务
# 值得注意的是，这里演示使用的ONNX模型是FP32精度的模型
# 如果你使用的fp16模型但传入的数据是fp32类型的会抛出数据异常的错误
# 另外ONNX的异常抛出是十分人性化的，它会指明你在推理是发生异常的具体位置以及应对策略
result = session.run(
                [output_name], {input_name: image.astype(np.float32)})[0]
result = result.argmax()

总结

以下ONNX简介来自于ONNX官方

ONNX 是一种用于表示机器学习模型的开放格式。 ONNX 定义了一组通用运算符——机器学习和深度学习模型的构建块——以及一种通用文件格式，使 AI 开发人员能够使用具有各种框架、工具、运行时和编译器的模型。

因此，ONNX是可以实现无缝的跨平台操作的。另外ONNX也支持了苹果的CoreML，这意味着如果你有需要在你的M1/M2 MacBook或者你的iOS设备上进行ONNX推理，ONNX也提供了对应支持的工具。

未来，ONNX将成为下一代AI研究人员或AI研发人员必备的技能之一。

到此这篇关于Python深度学习之Keras模型转换成ONNX模型流程详解的文章就介绍到这了,更多相关Python Keras模型转ONNX模型内容请搜索我们以前的文章或继续浏览下面的相关文章希望大家以后多多支持我们！