在 TensorFlow 中，当在 fit_generator 中使用 class_weights 时，会导致训练过程不断消耗越来越多的 CPU RAM，直至耗尽。每个时期之后内存使用量都会逐步增加。请参阅下面的可重现示例。为了保持可重现的示例较小，我减小了数据集的大小和批量大小，这显示了内存增加的趋势。在使用我的实际数据进行训练时，它耗尽了整个 128GB RAM 70 EPOCS。
有人遇到过这个问题或者对此有什么建议吗？我的数据有不平衡的数据，所以我必须使用 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()
