從Keras轉(zhuǎn)換成PB模型

請注意，如果直接使用Keras2ONNX進(jìn)行模型轉(zhuǎn)換大概率會出現(xiàn)報錯，這里筆者曾經(jīng)進(jìn)行過不同的嘗試，最后都失敗了。

所以筆者的推薦的情況是：首先將Keras模型轉(zhuǎn)換為TensorFlow PB模型。

那么通過tf.keras.models.load_model()這個函數(shù)將模型進(jìn)行加載，前提是你有一個基于h5格式或者hdf5格式的模型文件，最后再通過改變模型的保存格式save_format參數(shù)改為tf。即可實(shí)現(xiàn)從Keras轉(zhuǎn)換為TensorFow的格式文件了。

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

轉(zhuǎn)換后的模型文件結(jié)構(gòu)是這樣的：

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

2 directories, 4 files

從PB模型轉(zhuǎn)換成ONNX模型

從PB模型轉(zhuǎn)換為ONNX模型是很簡單的，通過調(diào)用tf2onnx這個模塊下的convert來進(jìn)行處理。

你需要做的，只需要將--save-model來指定你已經(jīng)轉(zhuǎn)換好的TensorFlow模型，使用--output來指定你的ONNX模型輸出的路徑（需要指定一個獨(dú)立的文件，如./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

改變現(xiàn)有的ONNX模型精度

考慮到在不同的計算設(shè)備上，半精度和雙精度鎖帶來的性能提升是顯而易見的。

這里我使用了一個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倍。

那么我們該如何將現(xiàn)有的ONNX模型從fp32模型轉(zhuǎn)換成fp16模型呢？

首先我們需要準(zhǔn)備一個叫onnxmltools的庫?？梢酝ㄟ^pip來進(jìn)行安裝。

pip install onnxmltools

確認(rèn)安裝好onnxmltools后，我們通過如下的一段腳本進(jìn)行精度的轉(zhuǎn)換：

import onnxmltools
# 加載float16_converter轉(zhuǎn)換器
from onnxmltools.utils.float16_converter import convert_float_to_float16
# 使用onnxmltools.load_model()函數(shù)來加載現(xiàn)有的onnx模型
# 但是請確保這個模型是一個fp32的原始模型
onnx_model = onnxmltools.load_model('../module/models/model.onnx')
# 使用convert_float_to_float16()函數(shù)將fp32模型轉(zhuǎn)換成半精度fp16
onnx_model_fp16 = convert_float_to_float16(onnx_model)
# 使用onnx.utils.save_model()函數(shù)來保存，
onnxmltools.utils.save_model(onnx_model_fp16, '../module/models/model_fp16.onnx')

部署ONNX 模型

在部署ONNX模型階段，我們將使用onnxruntime這個模塊。

針對你所將使用的計算設(shè)備，如果你是CPU用戶，那么你需要使用如下的指令來安裝onnxruntime

pip install onnxruntime

反之，如果你的計算設(shè)備是是GPU，那么你需要使用如下的指令來安裝onnxruntime

pip install onnxruntime-gpu

確認(rèn)好onnxruntime安裝完成后，你只需要使用如下的指令來加載你的ONNX模型即可

import onnxruntime as ort
# 指定onnx模型所在的位置
model_path = './module/models/model.onnx'
# 創(chuàng)建providers參數(shù)列表
providers = [
		# 指定模型可用的CUDA計算設(shè)備參數(shù)
        ('CUDAExecutionProvider', {
        	# 因?yàn)檫@里筆者只有一張GPU，因此GPU ID序列就為0
            'device_id': 0,
            # 這里網(wǎng)絡(luò)額外策略使用官方默認(rèn)值
            'arena_extend_strategy': 'kNextPowerOfTwo',
            # 官方這里默認(rèn)建議的GPU內(nèi)存迭代上限是2GB，如果你的GPU顯存足夠大
            # 可以將這里的2修改為其它數(shù)值
            'gpu_mem_limit': 2 * 1024 * 1024 * 1024,
            # cudnn轉(zhuǎn)換算法的調(diào)用參數(shù)設(shè)置為完整搜索
            'cudnn_conv_algo_search': 'EXHAUSTIVE',
            # 確認(rèn)從默認(rèn)流進(jìn)行CUDA流賦值
            'do_copy_in_default_stream': True,
        }),
        'CPUExecutionProvider',
    ]
# 使用onnxruntime.InferenceSession()函數(shù)創(chuàng)建Session
# 第一參數(shù)為模型所在的路徑，第二參數(shù)為模型的providers參數(shù)列表
session = ort.InferenceSession(model_path, providers=providers)
# 通過get_input()函數(shù)和get_output()函數(shù)獲取網(wǎng)絡(luò)的輸入和輸出名稱
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name 
# 使用session.run()函數(shù)執(zhí)行ONNX任務(wù)
# 值得注意的是，這里演示使用的ONNX模型是FP32精度的模型
# 如果你使用的fp16模型但傳入的數(shù)據(jù)是fp32類型的會拋出數(shù)據(jù)異常的錯誤
# 另外ONNX的異常拋出是十分人性化的，它會指明你在推理是發(fā)生異常的具體位置以及應(yīng)對策略
result = session.run(
                [output_name], {input_name: image.astype(np.float32)})[0]
result = result.argmax()

總結(jié)

以下ONNX簡介來自于ONNX官方

ONNX 是一種用于表示機(jī)器學(xué)習(xí)模型的開放格式。 ONNX 定義了一組通用運(yùn)算符——機(jī)器學(xué)習(xí)和深度學(xué)習(xí)模型的構(gòu)建塊——以及一種通用文件格式，使 AI 開發(fā)人員能夠使用具有各種框架、工具、運(yùn)行時和編譯器的模型。

因此，ONNX是可以實(shí)現(xiàn)無縫的跨平臺操作的。另外ONNX也支持了蘋果的CoreML，這意味著如果你有需要在你的M1/M2 MacBook或者你的iOS設(shè)備上進(jìn)行ONNX推理，ONNX也提供了對應(yīng)支持的工具。

未來，ONNX將成為下一代AI研究人員或AI研發(fā)人員必備的技能之一。

到此這篇關(guān)于Python深度學(xué)習(xí)之Keras模型轉(zhuǎn)換成ONNX模型流程詳解的文章就介紹到這了,更多相關(guān)Python Keras模型轉(zhuǎn)ONNX模型內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

最新評論