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

Question

在 TensorFlow 中，当在 fit_generator 中使用 class_weights 时，会导致训练过程不断消耗越来越多的 CPU RAM，直到耗尽。每个纪元后内存使用量都会逐步增加。请参阅下面的可重现示例。为了使可重现的示例变小，我减小了数据集的大小和批大小，这显示了内存增加的趋势。在使用我的实际数据进行训练时，它耗尽了 70 EPOCS 的全部 128GB RAM。

有人遇到过这个问题或者对此有什么建议吗？我的数据有不平衡的数据，所以我必须使用 class_weights 但我不能用它长时间运行训练。

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

第一张图片显示了使用 class_weights 的内存使用情况，而第二张图片显示了没有使用 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()

Answer 1

对于任何 future 的用户，每晚构建中似乎存在一个错误，该错误似乎在随后的每晚构建中得到修复。错误报告中有更多详细信息。

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