Python Keras模型转ONNX模型

时间:2022-09-29 FeathersMouch 人气:0

从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研发人员必备的技能之一。

加载全部内容