加快Tensorflow 2.0渐变胶带

时间:2019-03-03 11:58:38

标签: python tensorflow keras

我一直在跟踪位于here上的卷积VAE的TF 2.0教程。

由于急切,请手动计算渐变,然后使用tf.GradientTape()手动应用渐变。

for epoch in epochs:
  for x in x_train:
    with tf.GradientTape() as tape:
      loss = compute_loss(model, x)
    apply_gradients(tape.gradient(loss, model.trainable_variables))

该代码的问题在于它相当慢,每个纪元大约需要40-50秒。 如果我将批处理大小增加很多(到2048个左右),那么最终每个周期大约要花费8秒,但是模型的性能会下降很多。

另一方面,如果我使用一个更传统的模型(即使用基于懒惰图的模型而不是急切的模型),例如一个here,则即使在小批量。

model.add_loss(lazy_graph_loss)
model.fit(x_train epochs=epochs)

基于此信息,我的猜测是TF2.0代码的问题是损耗和梯度的手动计算。

有什么方法可以加快TF2.0代码的速度,使其更接近普通代码?

2 个答案:

答案 0 :(得分:3)

您的代码花费很长时间的主要原因是因为当我们使用没有任何张量的普通pythonic for循环时,图的构建会花费大量时间,因为从直觉上我们可能认为同一图用于训练的代码可能在每次迭代时都被重用,但是构造的图实际上是一个链状结构,其中每个节点都是训练子图,总编号为。该链中的节点的数量等于No。循环中的迭代。概括地说,张量流解开迭代,然后构造图。因此,仅在图形方面,就空间和时间而言,它就采用了大量冗余。太糟糕了,仅仅在一个正常的pythonic循环中重复添加两个张量(大约十亿次),就花费了将近半小时。

要解决此问题,特别是在您的情况下,我们可以在tf.data.Datasets api中使用.repeat转换的帮助,而不是编写

for i in range(epochs) :

我们可以写

For x in x_train.repeat(epochs) :
        Train here 

答案 1 :(得分:0)

我找到了解决方案:TensorFlow 2.0引入了functions的概念,将渴望的代码转换为图形代码。

用法非常简单。唯一需要做的更改是所有相关功能(如compute_lossapply_gradients)都必须用@tf.function注释。