Colab Pro+ CNN GPU memory saturation

Question

I have a problem when executing jupyter notebook for CNN in colab pro+, to train a model with a size of 560664x48x48x1. normally the data is composed of images with a size of 48x48. I used 10 fold cross-validation to train the model with the following code : rom sklearn.model_selection import StratifiedKFold

from datetime import datetime
from keras import backend as K
#here i will use stratifiedkfold cross validation method
##number of folds
kfold_splits=9
New_Data_train = np.reshape(New_Data_train , [560664,48, 48, 1])
skf = StratifiedKFold(n_splits=kfold_splits, shuffle=True)
lst_accu_stratified = []
lst_loss_stratified = []
lst_history_stratified = []
# K-fold Cross Validation model evaluation
fold_no = 1
#the folder where i wante to save the models 
FOLDER = "CNN_models_09_08_2022"
for k_train_index, k_val_index in skf.split(New_Data_train, New_Y_target):
    # Generate batches from indices
    xtrain, xval = New_Data_train[k_train_index], New_Data_train[k_val_index]
    ytrain, yval = New_Y_target[k_train_index], New_Y_target[k_val_index]
    ##ici je reinitialise le model temporaire (getting the model of cnn)
    model_CNN = get_model()
    #j'entraine le modele
    h = model_CNN.fit(xtrain,ytrain, epochs=100,batch_size=128)
    # evaluate the model
    scores = model_CNN.evaluate(xval,yval, verbose=0)
    ##I save the CNN model for future use.
    model_CNN.save(FOLDER+"/CNN_UNSW_"+str(fold_no)+"_FOLD.h5")
    lst_accu_stratified.append(scores[1])
    lst_loss_stratified.append(scores[0])
    lst_history_stratified.append(h)
    fold_no = fold_no+1
    del model_CNN
    #clearing sesssion
    K.clear_session()
    print("-------------------------------------------------------")

after executing the code above colab pro+ crashes because of insufficient GPU Ram and the following exception is raised :

InternalError: Failed copying input tensor from /job:localhost/replica:0/task:0/device:CPU:0 to /job:localhost/replica:0/task:0/device:GPU:0 in order to run _EagerConst: Dst tensor is not initialized.

can someone explain to me how to optimize this code in order to avoid this exception because it doesn't pass the 3 folds of cross validation. Thanks in advance

Answer 1

It is a known Keras problem that calling K.clear_session() and deleting the model sometimes fails to clear the GPU allocated memory, and this leads to the error you are seeing.

As a workaround, after clearing the session, you can force CUDA directly to release everything. For this, you can use numba. The code to release the CUDA memory with numba is as follows:

from numba import cuda
cuda.select_device(0)
cuda.close()