在 Tensorflow 2 中开始训练需要很长时间

Question

Tensorflow 2 需要大约 15 分钟来制作它的静态图（或者它在第一次通过之前所做的任何事情）。这之后的训练时间是正常的，但显然很难尝试等待任何反馈的 15 分钟。

生成器编码器和鉴别器是 Keras 模型中带有 GRU 单元的 RNN（未展开）。

生成器解码器是这样定义和调用的：

class GeneratorDecoder(tf.keras.layers.Layer):
def __init__(self, feature_dim):
    super(GeneratorDecoder, self).__init__()
    self.cell = tf.keras.layers.GRUCell(
        GRUI_DIM, activation='tanh', recurrent_activation='sigmoid',
        dropout=DROPOUT, recurrent_dropout=DROPOUT)
    self.batch_normalization = tf.keras.layers.BatchNormalization()
    self.dense = tf.keras.layers.Dense(
        feature_dim, activation='tanh')

@tf.function
def __call__(self, z, timesteps, training):
    # z has shape (batch_size, features)
    outputs = []
    output, state = z, z
    for i in range(timesteps):
        output, state = self.cell(inputs=output, states=state,
                                  training=training)
        dense_output = self.dense(
            self.batch_normalization(output))
        outputs.append(dense_output)
    return outputs

这是我的训练循环（mask_gt 和 missing_data 变量是使用 tf.cast 转换的，因此应该已经是张量）：

for it in tqdm(range(NO_ITERATIONS)):
   print(it)
   train_step()


@tf.function
def train_step():
    with tf.GradientTape(persistent=True) as tape:
        generator_output = generator(missing_data, training=True)
        imputed_data = get_imputed_data(missing_data, generator_output)
        mask_pred = discriminator(imputed_data)
        D_loss = discriminator.loss(mask_pred, mask_gt)
        G_loss = generator.loss(missing_data, mask_gt,
                                generator_output, mask_pred)
    gen_enc_grad = tape.gradient(
        G_loss, generator.encoder.trainable_variables)
    gen_dec_grad = tape.gradient(
        G_loss, generator.decoder.trainable_variables)
    disc_grad = tape.gradient(
        D_loss, discriminator.model.trainable_variables)
    del tape

    generator.optimizer.apply_gradients(
        zip(gen_enc_grad, generator.encoder.trainable_variables))
    generator.optimizer.apply_gradients(
        zip(gen_dec_grad, generator.decoder.trainable_variables))
    discriminator.optimizer.apply_gradients(
        zip(disc_grad, discriminator.model.trainable_variables))

注意，“0”是在几秒内打印出来的，所以慢的部分肯定不是更早的。这是调用的 get_imputed_data 函数：

def get_imputed_data(incomplete_series, generator_output):
    return tf.where(tf.math.is_nan(incomplete_series), generator_output, incomplete_series)

感谢您的回答！希望我提供的代码足以让您了解问题出在哪里。这是我阅读至少五年后第一次在这里发帖:)

我使用 Python 3.6 和 Tensorflow 2.1。

Answer 1

通过删除生成器和鉴别器调用函数的 tf.function 装饰器解决了这个问题。我在两个 tf.function 装饰函数中使用了一个全局 python 标量（迭代编号）。这导致每次都创建一个新图表（请参阅警告 in the tf.function docs）。

解决方案是删除使用的 python 变量或将它们转换为 tensorflow 变量。

在 Tensorflow 2 中开始训练需要很长时间

Starting training takes a very long time in Tensorflow 2

python

tensorflow

tensorflow2.0