要保存使用BatchNorm的TensorFlow模型,可以按照以下步骤进行操作:
import tensorflow as tf
model = tf.keras.Sequential([
tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(28, 28, 1)),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.MaxPooling2D((2, 2)),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
model.fit(train_images, train_labels, epochs=10)
model.save('model.h5')
这将保存模型的权重、架构和训练配置到名为model.h5
的文件中。
要加载已保存的模型,请使用以下代码:
loaded_model = tf.keras.models.load_model('model.h5')
这将加载模型的权重、架构和训练配置。你可以使用loaded_model
进行预测或继续训练。