tensorflow - Eager-Mode 非常慢(比 Graph-Mode 慢 22 倍)

Question

我读到 Tensorflow 2.0 将有一些重大变化，其中很大一部分将是 eager-execution [1]，所以我尝试玩了一下 tensorflow 的 eager-mode。

我从 github-repo 中获取了一段代码并尝试以 eager-mode 运行它(但是，没有按照建议使用 Keras-Model/Layers)。事实证明，它很慢。所以我尝试了不同的修改并将其与模型的原始来源(图形模式)进行了比较。结果是，graph-mode 比 eager-mode 快大约 22 倍。我很清楚，图形模式更快，但是这个数字？

情况总是如此，还是我需要对变量进行一些特殊修改/配置才能获得与图形模式相当的性能？

两种变体的源代码可以在 [2] 中找到。

提前致谢!

急切模式:

# With 
#  with tf.device("/gpu:0"):
#    ...
#
# Runtime is 0.35395
# Runtime is 0.12711
# Runtime is 0.12438
# Runtime is 0.12428
# Runtime is 0.12572
# Runtime is 0.12593
# Runtime is 0.12505
# Runtime is 0.12527
# Runtime is 0.12418
# Runtime is 0.12340

图形模式:

# Runtime is 0.81241
# Runtime is 0.00573
# Runtime is 0.00573
# Runtime is 0.00570
# Runtime is 0.00555
# Runtime is 0.00564
# Runtime is 0.00545
# Runtime is 0.00540
# Runtime is 0.00591
# Runtime is 0.00574

[1] https://groups.google.com/a/tensorflow.org/forum/#!topic/developers/JHDpgRyFVUs

[2] https://gist.github.com/lhlmgr/f6709e5aba4a5314b5221d58232b09bd

Answer 1

使用 Eager Execution 可能意味着取消使用 TensorFlow 图养成的一些习惯，因为过去运行一次的代码片段(例如，构建图以计算损失的 Python 函数)将重复运行(同一个 Python 函数现在将计算损失)在每次迭代中)。

我粗略地查看了提供的代码链接，并注意到一些简单的胜利，使用标准的 Python 分析工具也可能会看到这些胜利。您可能需要使用那些(cProfile、pyspy 等)

例如，Keras 网络目前实现为:

class NFModel(tf.keras.Model):
  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)

  def call(self, *args, **kwargs):
    num_layers = 6
    d, r = 2, 2
    bijectors = []

    for i in range(num_layers):
      with tf.variable_scope('bijector_%d' % i):
        V = tf.get_variable('V', [d, r], dtype=DTYPE)  # factor loading
        shift = tf.get_variable('shift', [d], dtype=DTYPE)  # affine shift
        L = tf.get_variable('L', [d * (d + 1) / 2], dtype=DTYPE)  # lower triangular
        bijectors.append(tfb.Affine(
          scale_tril=tfd.fill_triangular(L),
          scale_perturb_factor=V,
          shift=shift,
        ))

        alpha = tf.get_variable('alpha', [], dtype=DTYPE)
        abs_alpha = tf.abs(alpha) + .01
        bijectors.append(LeakyReLU(alpha=abs_alpha))

    base_dist = tfd.MultivariateNormalDiag(loc=tf.zeros([2], DTYPE))
    mlp_bijector = tfb.Chain(list(reversed(bijectors[:-1])), name='2d_mlp_bijector')
    dist = tfd.TransformedDistribution(distribution=base_dist, bijector=mlp_bijector)

相反，如果您在 __init__ 中创建一次变量并避免在每次调用网络时调用 tf.get_variable，您应该会看到很大的改进。

class NFModel(tf.keras.Model):
  def __init__(self, *args, **kwargs):
    super(NFModel, self).__init__(*args, **kwargs)
    num_layers = 6
    d, r = 2, 2
    self.num_layers = num_layers
    self.V = [tf.get_variable('V', [d, r], dtype=DTYPE)  for _ in range(num_layers)]
    self.shift = [tf.get_variable('shift', [d], dtype=DTYPE)   for _ in range(num_layers)]
    self.L = [tf.get_variable('L', [d * (d + 1) / 2], dtype=DTYPE)  for _ in range(num_layers)]
    self.alpha = [tf.get_variable('alpha', [], dtype=DTYPE) for _ in range(num_layers)]


  def call(self, *args, **kwargs):
    bijectors = []

    for i in range(self.num_layers):
      V = self.V[i]
      shift = self.shift[i]
      L = self.L[i]
      bijectors.append(tfb.Affine(
        scale_tril=tfd.fill_triangular(L),
        scale_perturb_factor=V,
        shift=shift,
      ))

      alpha = self.alpha[i]
      abs_alpha = tf.abs(alpha) + .01
      bijectors.append(LeakyReLU(alpha=abs_alpha))

    base_dist = tfd.MultivariateNormalDiag(loc=tf.zeros([2], DTYPE))
    mlp_bijector = tfb.Chain(list(reversed(bijectors[:-1])), name='2d_mlp_bijector')
    dist = tfd.TransformedDistribution(distribution=base_dist, bijector=mlp_bijector)

    return {"dist": dist}

可能还有其他类似的轻松胜利，一个分析工具会把你推向正确的方向。

另外，请注意，根据 RFC，TF 2.0 不是关于“急切执行”，而是更多关于如何与图形交互。

希望对您有所帮助。