To save a trained model in TensorFlow, follow these steps:
Create a
tf.keras.Model
object or subclass of it.Train the model using the
fit()
method.Create a
tf.keras.ModelCheckpoint
callback object and pass it to thefit()
method as an argument.Set the
save_weights_only
parameter toTrue
in theModelCheckpoint
callback object. This will save only the weights of the model, not the entire model structure.Set the
filepath
parameter in theModelCheckpoint
callback object to the desired file location where you want to save the model.Run the
fit()
method to train the model and save it at the specified file location.
Example:
Copy code# Create a model
model = tf.keras.Sequential([
tf.keras.layers.Dense(64, activation='relu', input_shape=(10,)),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax')
])
# Train the model
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
# Create a ModelCheckpoint callback to save the model weights
checkpoint = tf.keras.callbacks.ModelCheckpoint(filepath='/tmp/model.h5', save_weights_only=True)
# Train the model and save the weights
model.fit(x_train, y_train, epochs=10, callbacks=[checkpoint])
You can then load the saved model weights using the load_weights()
method of the model object:
Copy code# Load the saved model weights
model.load_weights('/tmp/model.h5')
# Evaluate the model on the test data
loss, accuracy = model.evaluate(x_test, y_test)
print(f'Test loss: {loss}, Test accuracy: {accuracy}')