Memory leak when using custom generator in Keras

Viewed 819

My problem is memory leak. Note:"GPU memory is no problem."

Environment:

Using Theano backend.
Keras==1.1.1, Theano==0.9.0 on Ubuntu 16.04.
16GB RAM and SWAP.
GPU: GTX1080

For the multiple outputs, the code at training was designed by oneself.
Look at the following function and training code.

Function:

def customGenerator(X,Y1,Y2):
    while True:
        idx = np.random.permutation(X.shape[0])
        datagen = ImageDataGenerator(rescale=1./255)
        batches = datagen.flow( X[idx], batch_size=32, shuffle=False)
        labels = [Y1[idx],Y2[idx]]
        idx0 = 0
        for batch in batches:
            idx1 = idx0 + batch.shape[0]
            yield batch, [labels[0][idx0:idx1],labels[1][idx0:idx1]]
            idx0 = idx1
            if idx1 >= X.shape[0]
                break

Training Code:

model.fit_generator(customGenerator(X_train,Y_train,Y_train2),
                    samples_per_epoch=X_train.shape[0]
                    nb_epoch=nb_epoch,
                    validation_data=customGenerator(X_val,Y_val,Y_val2),
                    nb_val_samples=X_val.shape[0],
                    verbose=1)

When training the network, it proceeded to 6 epochs and caused a memory error.
Observing the state of the memory showed that the memory usage rate increased as the epoch progressed.

I think there is a problem with the custom function, but I can not find where the problem is.

Could you help me?

0 Answers
Related