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

Question

Tensorflow 2 大约需要 15 分钟来制作静态图(或者在第一次传递之前它正在做什么(。之后的训练时间是正常的，但显然很难尝试等待任何反馈的 15 分钟。

生成器编码器和鉴别器是 RNN(未展开(，在 Keras 模型中具有 GRU 单元。

生成器解码器的定义和调用方式如下：

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 标量(迭代号(。这导致每次都会创建一个新图形(请参阅 tf.function 文档中的注意事项(。

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