keras转pb模型转onnx

首先先将h5模型转为pb模型

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import (
    Conv2D,
    MaxPooling2D,
    AveragePooling2D,
    Flatten,
    Dense,
    Dropout,
)
import tf2onnx
num_classes = 7
model = Sequential()

# 1st convolution layer
model.add(Conv2D(64, (5, 5), activation="relu", input_shape=(48, 48, 1)))
model.add(MaxPooling2D(pool_size=(5, 5), strides=(2, 2)))

# 2nd convolution layer
model.add(Conv2D(64, (3, 3), activation="relu"))
model.add(Conv2D(64, (3, 3), activation="relu"))
model.add(AveragePooling2D(pool_size=(3, 3), strides=(2, 2)))

# 3rd convolution layer
model.add(Conv2D(128, (3, 3), activation="relu"))
model.add(Conv2D(128, (3, 3), activation="relu"))
model.add(AveragePooling2D(pool_size=(3, 3), strides=(2, 2)))

model.add(Flatten())

# fully connected neural networks
model.add(Dense(1024, activation="relu"))
model.add(Dropout(0.2))
model.add(Dense(1024, activation="relu"))
model.add(Dropout(0.2))

model.add(Dense(num_classes, activation="softmax"))

model.load_weights("facial_expression_model_weights.h5")
model.save('tfmodel', save_format='tf',include_optimizer=False)

再用命令

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

然后我们尝试推理模型

import numpy as np

import onnxruntime as ort

# 加载ONNX模型

model_path = 'model.onnx'

session = ort.InferenceSession(model_path)

# 准备输入数据

input_name = session.get_inputs()[0].name

input_shape = session.get_inputs()[0].shape

# input_data = np.random.rand(*input_shape).astype(np.float32)

input_data = np.random.rand(1,48,48,1).astype(np.float32)

# 进行推理

output_name = session.get_outputs()[0].name

result = session.run([output_name], {
    
    input_name: input_data})[0]

print(result)

keras转pb模型转onnx

猜你喜欢