Whenever you train a model it takes a long time to train. Real-life models can take days or even weeks to train. If you cancel your training after it’s been running for a day or your model weights and values will be lost, then you would have to restart training from the beginning. But if you saved your model then you can always resume training from that point. Another benefit is that you can take your model and transfer to another computer, where you can continue training. In this tutorial, we’ll talk about saving and loading models.
Create a model
We use the CIFAR10 dataset to demonstrate model loading and saving. We normalize all pixel values to be between 0 and 1.
(x_train, y_train), (x_val, y_val) = tf.keras.datasets.cifar10.load_data() x_train = x_train.astype('float32') x_val = x_val.astype('float32') x_train /= 255 x_val /= 255 IMG_SIZE=32 BATCH_SIZE=32
Next is the model definition, which is defined in the
create_model function. This is a very basic model. We’re interested in learning how to load and save models, not creating the best model for the CIFAR10 dataset.
def create_model(): model = tf.keras.Sequential([ tf.keras.layers.Conv2D(kernel_size=3, filters=32, padding='same', activation='relu', input_shape=[IMG_SIZE,IMG_SIZE, 3]), tf.keras.layers.Conv2D(kernel_size=3, filters=64, padding='same', activation='relu'), tf.keras.layers.MaxPooling2D(pool_size=2), tf.keras.layers.Conv2D(kernel_size=3, filters=128, padding='same', activation='relu'), tf.keras.layers.Conv2D(kernel_size=1, filters=256, padding='same', activation='relu'), tf.keras.layers.GlobalAveragePooling2D(), tf.keras.layers.Dense(10,'softmax')]) model.compile(optimizer=tf.keras.optimizers.RMSprop(), loss=tf.keras.losses.SparseCategoricalCrossentropy(), metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]) return model
Save modal at the end of Epoch
We finally get to see how a model can be saved and perform model training by calling the fit method and providing keras callback.
ModelCheckpoint callback is used to save the model after each training epoch. It can save multiple files or a single file. Setting
save_weights_only to False in the Keras callback ‘ModelCheckpoint’ will save the full model.
Make sure to include the
epoch variable in your file path. Otherwise, your saved model will be replaced after every epoch. One option is to provide the
period parameter when creating the model checkpoint object. In this case 2, which as you can see saves a new model every 2 epochs.
history = model.fit(x_train[:2000], y_train[:2000], batch_size=BATCH_SIZE, epochs=6, callbacks=[model_checkpoint], validation_data=(x_val[:100], y_val[:100]), shuffle=True)
This will save a model every 2 epochs. If we look at the model directory, we can now see .h5 files. It contains the model’s configuration the model’s weights and the model’s optimizer’s state.
Save Final Model as HDF5 file
Another way of saving models is to call the
save() method on the model. This will create an HDF5 formatted file. The save method saves additional data, like the model’s configuration and even the state of the optimizer. A model that was saved using the save() method can be loaded with the function
This allows you to save the entirety of the state of a model in a single file.
Load Model and Continue training
The saved model can be re-instantiated in the exact same state, without any of the code used for model definition or training.
new_model = tf.keras.models.load_model('my_model.h5') new_model.evaluate(x_val,y_val)
The model returned by
load_model() is a compiled model ready to be used unless the saved model was not compiled. Re-compiling the model will reset the state of the model.
It is possible to save a partly train model and continue training after re-loading the model again. It’s useful when we have more training data in the future and we do not want to retrain the whole model again.
history = new_model.fit(x_train[:2000], y_train[:2000], batch_size=BATCH_SIZE, epochs=6, validation_data=(x_val[:100], y_val[:100]), shuffle=True)