A.I, Data and Software Engineering

Continue training big models on less powerful devices


It would not be a surprise that you may not have a powerful expensive machine to train a complicate model. You may experience the problem of not enough memory during training in some epoch. This article demonstrates a simple workaround for this.

The problem

Training deep learning models requires a lot of computing power. For most laptop and desktop today, you can still train the models but it can be a very time-consuming job. Unfortunately, for some model, your computer’ configuration is not enough to go through the task.

The common problem that you may encounter during training (e.g. Google Collab) is “Your session crashed after using all RAM “. What would be an inexpensive solution?

The workaround

Please note that this solution can ONLY be applied for machines which can finish loading data and complete at least one epoch of training.

Well, the idea is simple: save the training state to file, release all memory, reload the model and continue training the next epoch.

Create a model and check-points

Suppose that we have a deep learning model:

To save this model during training, we create a check-point callback using the following syntax:

As we want to save all state of the training model, we leave save_weights_only to False. Additionally, if we want to identify the saved model for each epoch, we can modify the file path with variables.

We use the callback during fitting the model.

The model will train and save the best model as follows.

Release memory and reload last saved model

if you use Jupyter Notebook, you can release the memory by using the following cell command.

To load the last saved model, we use tf.keras.models.load_model as shown in the following code:

Happy saving model and saving money for upgrading your device! Read more for save and load model in this article.

Add comment

A.I, Data and Software Engineering

PetaMinds focuses on developing the coolest topics in data science, A.I, and programming, and make them so digestible for everyone to learn and create amazing applications in a short time.

Pin It on Pinterest


You have successfully subscribed to the newsletter

There was an error while trying to send your request. Please try again.

Petaminds will use the information you provide on this form to be in touch with you and to provide updates.