I am trying to automate the (recursively) restart of a finished deep-learning training session in TensorFlow. Currently, to restart I am manually restarting my kernel and re-running the training code.


  1. I understand that "when training deep learning models, the model’s parameters, activations, and gradients are stored in the GPU memory." How would I clear the GPU memory without the need to manually restart my kernel?

  2. When I automate the restart of model training, do I need to restart from the very beginning (importing libraries + data preprocessing) OR can I just restart from where I start to build and fit the model?

  3. How would I implement this?

Thanks in advance!

Comment: This is how I call, compile, fit, and save the model.

# Get model 
def get_model():
    return build_model(input_shape, n_classes)

uNet_model = get_model() 

# Compile Model 

uNet_model.compile(optimizer= tf.optimizers.Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])

# Print Model Summary

# Fit Model 

# This is for a one-hot coded model: non-sparse 
history = uNet_model.fit(train_rgb_input, train_mask_categorical, 
                          validation_data=(val_rgb_input, val_mask_categorical), 
                          # class_weight=class_weights, 
                          verbose=1, shuffle=True) 

# Save model

1 Answer 1


What do you mean by restart training?

  1. Re-train your model? Then you just need to rebuild your model (and delete the old model to free up GPU memory).

  2. Continue training If you just want to continue training, e.g. because you realized that the model needs more epochs, you can just call the .fit method again.

A general remark: There is often a risk of running out of GPU memory. TF should release automatically allocated memory nowadays (e.g. batches of your training data), but if you have for example several models active at the same time, you can still run out of memory. You can deal with this later by storing the weights as numpy arrays (or just saving them to disk). Then you can load them back into the model when you need them.

  • $\begingroup$ I mean to re-train the same model, again, and not to extend that training session for more epochs. $\endgroup$ Jan 22 at 14:26
  • $\begingroup$ You could save the initial weights at the beginning and load them into your model after each run. $\endgroup$
    – Max
    Jan 22 at 14:34
  • $\begingroup$ Thank you for your comment. I guess the next question is will your proposed method release/erase/relieve the previously stored data in the Nvidia GPU? My goal is to use an "empty" (don't know if this is the correct technical term) GPU and re-train the model recursively. $\endgroup$ Jan 22 at 15:37

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.