TensorFlow2.x冻结模型，保存为.pb格式方便部署

最新推荐文章于 2024-08-24 10:07:55 发布

布布汪

最新推荐文章于 2024-08-24 10:07:55 发布

阅读量8.4k

点赞数 12

文章标签： tensorflow 深度学习机器学习神经网络

本文链接：https://blog.csdn.net/huangcong159/article/details/106882311

版权

本文介绍了如何使用TensorFlow2.x创建并训练模型，然后将模型保存为.h5格式，接着通过定义转化函数将模型冻结为.pb格式，以实现跨平台部署，并且这种方式不受TensorFlow版本限制。

摘要由CSDN通过智能技术生成

1、用tf.keras创建模型

inputs = tf.keras.Input(shape=(1000, 9,), name='input')
input = tf.keras.layers.Flatten(name="flatten")(inputs)
f1 = tf.keras.layers.Dense(512, activation='relu', name='dense_1')(input)
d1 = tf.keras.layers.Dropout(0.2,name='dropout_1')(f1)
f2 = tf.keras.layers.Dense(128, activation='relu', name='dense_2')(d1)
outputs = tf.keras.layers.Dense(3, activation='softmax', name='output')(f2)
model = tf.keras.Model(inputs=inputs, outputs=outputs, name='intrusion')
model.summary()

2、模型训练

model.compile(loss='sparse_categorical_crossentropy',
              optimizer=tf.keras.optimizers.Adam(),
              metrics=['sparse_categorical_accuracy'])
history = model.fit(x_train, y_train,