带有 Tensorflow (1.3) 后端的 Keras (2.0.8) 占用所有可用内存
Keras (2.0.8) with Tensorflow (1.3) backend takes all available RAM
我正在使用 keras 库，启用了 tensorflow 后端和 CUDA。查看 PIP 包版本输出：
Keras (2.0.8)
tensorflow-gpu (1.3.0)
tensorflow-tensorboard (0.1.8)
我有以下代码创建 VGG16 模型并加载 ImageNet 权重：
def create_vgg16_model(target_size: tuple, n_classes: int):
base = VGG16(include_top=False,
input_shape=target_size,
weights='imagenet')
x = base.output
x = Flatten()(x)
x = Dense(n_classes, activation='softmax', name='top')(x)
model = Model(inputs=base.input, outputs=x)
for layer in model.layers[:-1]:
layer.trainable = False
model.compile(optimizer='adam', loss='categorical_crossentropy')
return model
模型的训练进展顺利，nvidia-smi 显示 GPU 内存已按需使用。但是后来我检查了 top 命令的输出，这是我看到的：
PID USER PR NI VIRT RES SHR S %CPU %MEM TIME+ COMMAND
1268 ck 20 0 166288 31964 12416 S 29.5 0.1 13:05.39 Xtightvnc
32235 ck 30 10 32252 3700 3356 S 5.3 0.0 0:36.93 cwaves
------------------------------------------------------------------------------
32212 ck 20 0 27.485g 1.184g 190404 S 2.3 3.8 0:35.44 python
------------------------------------------------------------------------------
26015 root 20 0 0 0 0 S 0.3 0.0 0:00.30 kworker/3:1
31754 ck 20 0 43168 3904 3080 R 0.3 0.0 0:04.45 top
1 root 20 0 185644 6204 3984 S 0.0 0.0 0:10.44 systemd
我已经使用调试器遍历了代码，并意识到内存是在从 keras.backend.tensorflow_backed 获取的以下函数中分配的，它创建了一个 tf.Session 对象：
def get_session():
global _SESSION
if tf.get_default_session() is not None:
session = tf.get_default_session()
else:
if _SESSION is None:
if not os.environ.get('OMP_NUM_THREADS'):
config = tf.ConfigProto(allow_soft_placement=True)
else:
num_thread = int(os.environ.get('OMP_NUM_THREADS'))
config = tf.ConfigProto(intra_op_parallelism_threads=num_thread,
allow_soft_placement=True)
# next line allocates ~28GB of RAM
_SESSION = tf.Session(config=config)
session = _SESSION
if not _MANUAL_VAR_INIT:
with session.graph.as_default():
_initialize_variables()
return session
而且，所有可用模型都会发生这种情况，因为内存是在会话创建时、训练开始或变量初始化之前分配的。
我知道 TF 会分配所有可用的 GPU 内存（除非你覆盖 ConfigProto and/or 调整你的环境变量），但它对 RAM 做同样的事情吗？ IE。似乎框架正在分配我机器上的所有 RAM，但已经由其他进程分配的内存除外。
有没有人在不同版本的 tensorflow 或 keras 中发现这种行为？您认为有什么方法可以限制使用的内存量吗？
更新 1
前段时间，我的一个训练脚本在 50-60 个训练阶段后被内核杀死，出现内存不足错误。尽管 volatile GPU 内存使用统计显示它也被使用。 （不只是分配，据我所知）。
更新 2
同意，虚拟内存不是有效的指标。但我发现在模型训练过程中内存消耗几乎呈线性增长。我有以下训练循环：
def train_model(model, x, y):
loss = model.train_on_batch(x, y)
return loss
def train_model_42(model, x, y):
# dummy function
return 42.0
def training_loop():
# training parameters
target_size = 224, 224, 3
batch_size = 128
# generator yielding batches of file paths
files_stream = FilesStream(folder=TRAIN_IMAGES, batch_size=batch_size)
files_source = files_stream()
# list of generators loading images from persistent storage
gens = [
image_loader(),
augment_images(horizontal_flip=True),
shuffle_samples(),
normalize_images(target_size=target_size)
]
# Model from keras.applications with replaced top layer
model = get_model('resnet50').build(n_classes=n_classes)
for epoch in range(1, 1001):
epoch_loss = []
for _ in range(files_stream.steps_per_epoch):
for gen in gens:
gen.send(None)
processed = next(files_source)
for gen in gens:
processed = gen.send(processed)
x, y = processed
loss = train_model_42(model, x, y) # <-- this shows pic. 1
# loss = train_model(model, x, y) <-- this shows pic. 2
epoch_loss.append(loss)
avg_loss = sum(epoch_loss) / len(epoch_loss)
print('Epoch %03d: train loss = %2.4f' % (epoch, avg_loss))
当我使用虚拟训练函数时，内存消耗图看起来像 pic 1 上显示的那样：
但是 运行 一个真正的训练过程，它看起来像 pic 2：
为什么训练过程中内存消耗越来越大？是否缓存了以前的批次数据？ model/weights 或其他任何东西是否应该占用越来越多的内存？
我认为我的数据预处理管道可能有问题，但我有意将预处理函数编写为生成器。会不会是某种默认的 Keras 回调应用于跟踪训练信息的模型，这会增加内存使用量？
我想我找到了问题的根源。果然，它与tensorflow或keras无关，而是我使用它们的方法。
这是一个类似于我的图像增强函数的函数：
def augment_images():
transformer = ImageDataGenerator()
while True:
x, y = yield
generator = transformer.flow(x, y, batch_size=len(x), shuffle=False)
transformed = next(generator)
yield transformed
它使用 keras.preprocessing.image.ImageDataGenerator class 来增强图像。但是 class 本身实例化 NumpyArrayIterator 对象，该对象 保持对 x 和 y 的引用 并调用 ImageDataGenerator 作为委托.而且，这就是内存泄漏的根源。似乎这些对象阻止了数组被垃圾回收。
这是一个明确使用迭代器的更新增强函数：
def augment_images(width_shift=0.2,
height_shift=0.2,
zoom=0.2,
rotation=30,
vertical_flip=False,
horizontal_flip=False):
transformer = ImageDataGenerator()
iterator = None
while True:
x, y = yield
if iterator is None:
iterator = NumpyArrayIterator(
x, y, transformer,
batch_size=len(x),
shuffle=False,
seed=None,
data_format=transformer.data_format)
else:
iterator.n = x.shape[0]
iterator.x = x
iterator.y = y
transformed = next(iterator)
yield transformed
所以，问题出在我用来预处理数据的生成器包装器中。 （或者我会说，在我使用 Keras 的 API 和 Python 的生成器的方法中）。至少现在，当我替换图像增强功能时，不再有内存泄漏。
我正在使用 keras 库，启用了 tensorflow 后端和 CUDA。查看 PIP 包版本输出：
Keras (2.0.8)
tensorflow-gpu (1.3.0)
tensorflow-tensorboard (0.1.8)
我有以下代码创建 VGG16 模型并加载 ImageNet 权重：
def create_vgg16_model(target_size: tuple, n_classes: int):
base = VGG16(include_top=False,
input_shape=target_size,
weights='imagenet')
x = base.output
x = Flatten()(x)
x = Dense(n_classes, activation='softmax', name='top')(x)
model = Model(inputs=base.input, outputs=x)
for layer in model.layers[:-1]:
layer.trainable = False
model.compile(optimizer='adam', loss='categorical_crossentropy')
return model
模型的训练进展顺利，nvidia-smi 显示 GPU 内存已按需使用。但是后来我检查了 top 命令的输出，这是我看到的：
PID USER PR NI VIRT RES SHR S %CPU %MEM TIME+ COMMAND
1268 ck 20 0 166288 31964 12416 S 29.5 0.1 13:05.39 Xtightvnc
32235 ck 30 10 32252 3700 3356 S 5.3 0.0 0:36.93 cwaves
------------------------------------------------------------------------------
32212 ck 20 0 27.485g 1.184g 190404 S 2.3 3.8 0:35.44 python
------------------------------------------------------------------------------
26015 root 20 0 0 0 0 S 0.3 0.0 0:00.30 kworker/3:1
31754 ck 20 0 43168 3904 3080 R 0.3 0.0 0:04.45 top
1 root 20 0 185644 6204 3984 S 0.0 0.0 0:10.44 systemd
我已经使用调试器遍历了代码，并意识到内存是在从 keras.backend.tensorflow_backed 获取的以下函数中分配的，它创建了一个 tf.Session 对象：
def get_session():
global _SESSION
if tf.get_default_session() is not None:
session = tf.get_default_session()
else:
if _SESSION is None:
if not os.environ.get('OMP_NUM_THREADS'):
config = tf.ConfigProto(allow_soft_placement=True)
else:
num_thread = int(os.environ.get('OMP_NUM_THREADS'))
config = tf.ConfigProto(intra_op_parallelism_threads=num_thread,
allow_soft_placement=True)
# next line allocates ~28GB of RAM
_SESSION = tf.Session(config=config)
session = _SESSION
if not _MANUAL_VAR_INIT:
with session.graph.as_default():
_initialize_variables()
return session
而且，所有可用模型都会发生这种情况，因为内存是在会话创建时、训练开始或变量初始化之前分配的。
我知道 TF 会分配所有可用的 GPU 内存（除非你覆盖 ConfigProto and/or 调整你的环境变量），但它对 RAM 做同样的事情吗？ IE。似乎框架正在分配我机器上的所有 RAM，但已经由其他进程分配的内存除外。
有没有人在不同版本的 tensorflow 或 keras 中发现这种行为？您认为有什么方法可以限制使用的内存量吗？
更新 1
前段时间，我的一个训练脚本在 50-60 个训练阶段后被内核杀死，出现内存不足错误。尽管 volatile GPU 内存使用统计显示它也被使用。 （不只是分配，据我所知）。
更新 2
同意，虚拟内存不是有效的指标。但我发现在模型训练过程中内存消耗几乎呈线性增长。我有以下训练循环：
def train_model(model, x, y):
loss = model.train_on_batch(x, y)
return loss
def train_model_42(model, x, y):
# dummy function
return 42.0
def training_loop():
# training parameters
target_size = 224, 224, 3
batch_size = 128
# generator yielding batches of file paths
files_stream = FilesStream(folder=TRAIN_IMAGES, batch_size=batch_size)
files_source = files_stream()
# list of generators loading images from persistent storage
gens = [
image_loader(),
augment_images(horizontal_flip=True),
shuffle_samples(),
normalize_images(target_size=target_size)
]
# Model from keras.applications with replaced top layer
model = get_model('resnet50').build(n_classes=n_classes)
for epoch in range(1, 1001):
epoch_loss = []
for _ in range(files_stream.steps_per_epoch):
for gen in gens:
gen.send(None)
processed = next(files_source)
for gen in gens:
processed = gen.send(processed)
x, y = processed
loss = train_model_42(model, x, y) # <-- this shows pic. 1
# loss = train_model(model, x, y) <-- this shows pic. 2
epoch_loss.append(loss)
avg_loss = sum(epoch_loss) / len(epoch_loss)
print('Epoch %03d: train loss = %2.4f' % (epoch, avg_loss))
当我使用虚拟训练函数时，内存消耗图看起来像 pic 1 上显示的那样：
但是 运行 一个真正的训练过程，它看起来像 pic 2：
为什么训练过程中内存消耗越来越大？是否缓存了以前的批次数据？ model/weights 或其他任何东西是否应该占用越来越多的内存？
我认为我的数据预处理管道可能有问题，但我有意将预处理函数编写为生成器。会不会是某种默认的 Keras 回调应用于跟踪训练信息的模型，这会增加内存使用量？
我想我找到了问题的根源。果然，它与tensorflow或keras无关，而是我使用它们的方法。
这是一个类似于我的图像增强函数的函数：
def augment_images():
transformer = ImageDataGenerator()
while True:
x, y = yield
generator = transformer.flow(x, y, batch_size=len(x), shuffle=False)
transformed = next(generator)
yield transformed
它使用 keras.preprocessing.image.ImageDataGenerator class 来增强图像。但是 class 本身实例化 NumpyArrayIterator 对象，该对象 保持对 x 和 y 的引用 并调用 ImageDataGenerator 作为委托.而且，这就是内存泄漏的根源。似乎这些对象阻止了数组被垃圾回收。
这是一个明确使用迭代器的更新增强函数：
def augment_images(width_shift=0.2,
height_shift=0.2,
zoom=0.2,
rotation=30,
vertical_flip=False,
horizontal_flip=False):
transformer = ImageDataGenerator()
iterator = None
while True:
x, y = yield
if iterator is None:
iterator = NumpyArrayIterator(
x, y, transformer,
batch_size=len(x),
shuffle=False,
seed=None,
data_format=transformer.data_format)
else:
iterator.n = x.shape[0]
iterator.x = x
iterator.y = y
transformed = next(iterator)
yield transformed
所以，问题出在我用来预处理数据的生成器包装器中。 （或者我会说，在我使用 Keras 的 API 和 Python 的生成器的方法中）。至少现在，当我替换图像增强功能时，不再有内存泄漏。