TensorFlow-在fit_generator中使用class_weights会导致内存泄漏

时间:2019-08-02 00:56:17

标签: tensorflow keras

在TensorFlow中,在fit_generator中使用class_weights时,导致训练过程持续消耗越来越多的CPU RAM,直到耗尽为止。在每个时期之后,内存使用量都会逐步增加。请参见下面的可复制示例。为了使可重现的示例保持较小,我减小了数据集的大小和批处理大小,这表明内存增加的趋势。在训练我的实际数据时,它消耗了70个EPOCS的全部128GB RAM。

有人遇到这个问题或对此有任何建议吗?我的数据具有不平衡的数据,因此我必须使用class_weights,但是我不能为此长时间进行训练。

在下面的代码示例中,如果您注释掉了类权重,则程序将在不耗尽内存的情况下进行训练。

第一张图片显示了使用class_weights的内存使用情况,第二张图片显示了不使用class_weights的使用情况。

Memory Usage with class_weights

Memory usage without class_weights

import tensorflow as tf
tf.enable_eager_execution()
import numpy as np

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import CuDNNLSTM, Dense
from tensorflow.keras.optimizers import Adadelta


feature_count = 25
batch_size = 16
look_back = 5
target_groups = 10

def random_data_generator( ):
    x_data_size =(batch_size, look_back, feature_count) # batches, lookback, features
    x_data = np.random.uniform(low=-1.0, high=5, size=x_data_size)

    y_data_size = (batch_size, target_groups)
    Y_data = np.random.randint(low=1, high=21, size=y_data_size)

    return x_data, Y_data

def get_simple_Dataset_generator():        
    while True:
        yield random_data_generator()

def build_model():
    model = Sequential()
    model.add(CuDNNLSTM(feature_count,
                    batch_input_shape=(batch_size,look_back, feature_count),
                    stateful=False))  
    model.add(Dense(target_groups, activation='softmax'))
    optimizer = Adadelta(learning_rate=1.0, epsilon=None) 
    model.compile(loss='categorical_crossentropy', optimizer=optimizer) 
    return model


def run_training():

    model = build_model()
    train_generator = get_simple_Dataset_generator()
    validation_generator = get_simple_Dataset_generator()
    class_weights = {0:2, 1:8, 2:1, 3:4, 4:8, 5:35, 6:30, 7:4, 8:5, 9:3}

    model.fit_generator(generator = train_generator,
            steps_per_epoch=1,
            epochs=1000,            
            verbose=2,
            validation_data=validation_generator,
            validation_steps=20,
            max_queue_size = 10,
            workers = 0, 
            use_multiprocessing = False,
            class_weight = class_weights
            )

if __name__ == '__main__': 
    run_training()

1 个答案:

答案 0 :(得分:1)

对于任何将来的用户,夜间版本中似乎都有一个错误,此错误似乎在以后的夜间版本中已修复。错误报告中的更多详细信息。

https://github.com/tensorflow/tensorflow/issues/31253