python - 在tensorflow的keras的fit_generator中，AsyncResult在意外情况下挂起

Question

这是 issue I posted on the tensorflow Github 的复制粘贴。

系统信息

• 我是否编写了自定义代码:是
• 操作系统平台和发行版:Linux Ubuntu 16.04
• TensorFlow 安装自:pip
• TensorFlow 版本:2.0.0b1
• Python版本:3.6.8
• CUDA/cuDNN 版本:V10.0.130
• GPU 型号和内存:Quadro P5000 (16GB)

描述当前行为

我有一个非常复杂的模型来解决图像到图像的问题。我还使用了一个自定义回调，它有时会使用 numpy 生成一些噪音。当我在这个模型上使用 fit_generator 时，它成功执行第一个纪元，然后在第二个、第三个或第四个纪元上卡在纪元的开头。我设法看到问题发生在哪里，它发生在这里:https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/utils/data_utils.py#L875基本上，如果我在第二个 get 上设置超时，它会在几个成功的时期(有时只有一个)后超时。没有抛出任何错误，所以我不知道它为什么挂起。此外，如果我在代码中进行调试，我可以同步执行该函数，一切都会正常工作。

重现问题的代码

我没有设法使用 fit_generator 获得一个最小的示例(基本上它过于依赖我使用我的复杂模型)。但是，我有一个最小的示例，当我模仿 model_iteration 函数时，它会重现该错误。您需要安装以下内容才能使其工作:pip install tensorflow-gpu==2.0.0b1 numpy tqdm

# imports
import time

import numpy as np
import tensorflow as tf
from tensorflow.python.keras import callbacks as cbks
from tensorflow.keras.callbacks import Callback
from tensorflow.python.keras.models import Model
from tensorflow.python.keras.engine import training_utils
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.python.keras.utils import data_utils
from tensorflow.python.keras.utils import generic_utils
from tqdm import tqdm_notebook

# helper function (taken from https://github.com/tensorflow/tensorflow/blob/r2.0/tensorflow/python/keras/engine/training_generator.py#L500)
def _make_enqueued_generator(generator,
                             workers=1,
                              use_multiprocessing=False,
                             max_queue_size=10,
                             shuffle=False):    
    enqueuer = data_utils.OrderedEnqueuer(
        generator, use_multiprocessing=use_multiprocessing, shuffle=shuffle)
    enqueuer.start(workers=workers, max_queue_size=max_queue_size)
    output_generator = enqueuer.get()
    return output_generator, enqueuer

# My silly callback
class Noise(Callback):
     def on_batch_end(self, batch, logs={}):
        image_shape = [1, 2**7, 2**7, 1]
        noise = np.random.normal(scale=1.0, size=image_shape)

# My data
batch_size = 8
n_samples_train = 720
x = np.random.rand(n_samples_train, 256, 256, 1)
im_gen_train = ImageDataGenerator().flow(x, batch_size=batch_size)


# My training set up (to mimic https://github.com/tensorflow/tensorflow/blob/r2.0/tensorflow/python/keras/engine/training_generator.py#L41)
data = im_gen_train
steps_per_epoch = int(n_samples_train / batch_size)
epochs = 20
max_queue_size=35
workers=35
use_multiprocessing=True
shuffle=False
initial_epoch=0
mode=1
steps_name='steps'
noise_cb = Noise()
noise_cb.on_train_batch_end = noise_cb.on_batch_end
callbacks=[noise_cb]

generator, enqueuer = _make_enqueued_generator(
    im_gen_train,
    workers=workers,
    use_multiprocessing=use_multiprocessing,
    max_queue_size=max_queue_size,
    shuffle=shuffle)

callbacks = cbks.configure_callbacks(
    callbacks,
    Model(),
    do_validation=False,
    epochs=epochs,
    steps_per_epoch=steps_per_epoch,
    batch_size=batch_size,
    samples=n_samples_train,
    verbose=0,  # Handle ProgBar as part of Callbacks once hooks are ready.
    mode=mode,
)
callbacks._call_begin_hook(mode)

for epoch in tqdm_notebook(range(initial_epoch, epochs)):
    callbacks.on_epoch_begin(epoch, {})

    for step in tqdm_notebook(range(steps_per_epoch), leave=False):
        callbacks._call_batch_hook('train', 'begin', step, {})
        batch_data = next(generator)

        # I don't actually train a model, so I just sleep for this time, this would be the backprop
        time.sleep(0.1)
        callbacks._call_batch_hook('train', 'end', step, {})

如果您保持原样，它将在大约 1、2、3 或 4 次迭代后挂起。您可以注释掉 noise = np.random.normal(scale=1.0, size=image_shape) 行，看看它没有挂起。

您还可以在第二个get中修改tensorflow的源代码和超时here，以便进行调试。

另请注意，如果 sleep 时间不够长，则不会出现悬挂。

编辑

我终于成功地整理了一个直接涉及 fit_generator 的最小示例:

# imports
import time

from keras_tqdm import TQDMNotebookCallback
import numpy as np
import tensorflow as tf
from tensorflow.keras.callbacks import Callback
from tensorflow.keras.layers import Input, Conv2D, Lambda, concatenate
from tensorflow.python.keras.models import Model
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.utils import  Sequence


# My silly callback
class Noise(Callback):
     def on_batch_end(self, batch, logs={}):
        image_shape = [1, 2**7, 2**7, 1]
        noise = np.random.normal(scale=1.0, size=image_shape)

# my metrics
def keras_psnr(y_true, y_pred):
    max_pixel = tf.math.reduce_max(y_true)
    min_pixel = tf.math.reduce_min(y_true)
    return tf.image.psnr(y_true, y_pred, max_pixel - min_pixel)

def keras_ssim(y_true, y_pred):
    max_pixel = tf.math.reduce_max(y_true)
    min_pixel = tf.math.reduce_min(y_true)
    return tf.image.ssim(y_true, y_pred, max_pixel - min_pixel)

# My data
class MergedGenerators(Sequence):
    def __init__(self, *generators):
        self.generators = generators
        # TODO add a check to verify that all generators have the same length

    def __len__(self):
        return len(self.generators[0])

    def __getitem__(self, index):
        return tuple([generator[index] for generator in self.generators])

batch_size = 8
n_samples_train = 720
size = 256
x = np.random.rand(n_samples_train, size, size, 1)
im_gen_train_1 = ImageDataGenerator().flow(x, batch_size=batch_size, seed=0)
im_gen_train_2 = ImageDataGenerator().flow(x, batch_size=batch_size, seed=0)
im_gen_train = MergedGenerators(im_gen_train_1, im_gen_train_2)

# my fake model
im = Input((None, None, 1))
conv = Conv2D(256, 3, padding='same')(im)
conv = Conv2D(256, 3, padding='same')(conv)
conv = Conv2D(1, 3, padding='same')(conv)
ident = Lambda(lambda x: x)(conv)
model = Model(im, ident)
model.compile(loss='mse', optimizer='adam', metrics=[keras_psnr, keras_ssim])
print(model.summary(line_length=150))

# My training set up
noise_cb = Noise()
noise_cb.on_train_batch_end = noise_cb.on_batch_end
tqdm_cb = TQDMNotebookCallback(metric_format="{name}: {value:e}")
tqdm_cb.on_train_batch_begin = tqdm_cb.on_batch_begin
tqdm_cb.on_train_batch_end = tqdm_cb.on_batch_end
model.fit_generator(
    im_gen_train,
    steps_per_epoch=int(n_samples_train / batch_size), 
    epochs=20,
    max_queue_size=35,
    workers=35,
    use_multiprocessing=True,
    shuffle=False,
    callbacks=[noise_cb, tqdm_cb],
    verbose=0,
)

它不是很简单，但至少可以直接使用(你只需要另外安装keras-tqdm)。

Answer 1

此问题已在 2.1 版本中解决。

另一个修复方法是使用 new random number generation API of numpy按照建议here 。这会将行 noise = np.random.normal(scale=1.0, size=image_shape) 更改为 noise = np.random.default_rng().normal(scale=1.0, size=image_shape) )。此修复甚至在 2.0 版本中也有效。

这是我在 Github 上给出的答案的复制粘贴。