tf.experimental.set_loop_options - 形状不变量因 TPU 策略而失败

Question

我试图在 tensorflow 的分布式策略中运行一个循环，其中变量不断存储输出，因此它的形状不断变化。

import tensorflow as tf
strategy = tf.distribute.MirroredStrategy()

@tf.function(autograph=True)
def f():
  v = tf.constant((0,))
  for i in tf.range(3):
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(v, tf.TensorShape([None]))]
    )
    v = tf.concat((v, [i]), 0)
  return v

strategy.run(f)

上面的代码工作正常。但是如果我用 TPUstrategy 替换策略，即下面的代码

import tensorflow as tf
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

@tf.function(autograph=True)
def f():
  v = tf.constant((0,))
  for i in tf.range(3):
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(v, tf.TensorShape([None]))]
    )
    v = tf.concat((v, [i]), 0)
  return v

strategy.run(f)

此代码引发错误 -

InvalidArgumentError                      Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/IPython/core/formatters.py in __call__(self, obj)
    697                 type_pprinters=self.type_printers,
    698                 deferred_pprinters=self.deferred_printers)
--> 699             printer.pretty(obj)
    700             printer.flush()
    701             return stream.getvalue()

6 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py in shape(self)
   1325         # `_tensor_shape` is declared and defined in the definition of
   1326         # `EagerTensor`, in C.
-> 1327         self._tensor_shape = tensor_shape.TensorShape(self._shape_tuple())
   1328       except core._NotOkStatusException as e:
   1329         raise core._status_to_exception(e) from None

InvalidArgumentError: 9 root error(s) found.
  (0) INVALID_ARGUMENT: {{function_node __inference_tpu_function_230}} Input and output shapes of loop body do not match: (s32[], s32[], s32[], s32[1], s32[], /*index=5*/s32[]) vs. (s32[], s32[], s32[], s32[2], s32[], /*index=5*/s32[])

     [[{{function_node __inference_f_202}}{{node while}}]]
     [[TPUReplicate/_compile/_17449455679645184466/_2]]
  (1) INVALID_ARGUMENT: {{function_node __inference_tpu_function_230}} Input and output shapes of loop body do not match: (s32[], s32[], s32[], s32[1], s32[], /*index=5*/s32[]) vs. (s32[], s32[], s32[], s32[2], s32[], /*index=5*/s32[])

     [[{{function_node __inference_f_202}}{{node while}}]]
     [[TPUReplicate/_compile/_17449455679645184466/_2]]
     [[tpu_compile_succeeded_assert/_5804319525153755267/_3/_35]]
  (2) INVALID_ARGUMENT: {{function_node __inference_tpu_function_230}} Input and output shapes of loop body do not match: (s32[], s32[], s32[], s32[1], s32[], /*index=5*/s32[]) vs. (s32[], s32[], s32[], s32[2], s32[], /*index=5*/s32[])

     [[{{function_node __inference_f_202}}{{node while}}]]
     [[TPUReplicate/_compile/_17449455679645184466/_2]]
     [[cluster_tpu_function/control_after/_1/_99]]
  (3) INVALID_ARGUMENT: {{function_node __inference_tpu_function_230}} Input and output shapes of loop body do not match: (s32[], s32[], s32[], s32[1], s32[], /*index=5*/s32[]) vs. (s32[], s32[], s32[], s32[2], s32[], /*index=5*/s32[])

     [[{{function_node __inference_f_202}}{{node while}}]]
     [[TPUReplicate/_compile/_17449455679645184466/_2]]
     [[cluster_tpu_function/control_after/_1/_95]]
  (4) INVALID_ARGUMENT: {{function_node __inference_tpu_function_230}} Input and output shapes of loop body do not match: (s32[], s32[], s32[], s32[1], s32[], /*index=5*/s32[]) vs. (s32[], s32[], s32[], s32[2], s32[], /*index=5*/s32[])

     [[{{function_node __inference_f_202}}{{node while}}]]
     [[TPUReplicate/_compile/_17449455679645184466/_2]]
     [[cluster_tpu_function/control_after/_1/_91]]
  (5) INVALID_ARGUMENT: {{function_node __inference_tpu_function_230}} Input and output shapes of loop body do not match: (s32[], s32[], s32[], s32[1], s32[], /*index=5*/s32[]) vs. (s32[], s32[], s32[], s32[2], s32[], /*index=5*/s32[])

     [[{{function_node __inference_f_202}}{{node while}}]]
     [[TPUReplicate/_compile/_17449455679645184466/_2]]
     [[cluster_tpu_function/control_after/_1/_83]]
  (6) INVALID_ARGUMENT: {{function_node __inference_tpu_function_230}} Input and output shapes of loop body do not match: (s32[], s32[], s32[], s32[1], s32[], /*index=5*/s32[]) vs. (s32[], s32[], s32[], s32[2], s32[], /*index=5*/s32[])

     [[{{function_node __inference_f_202}}{{node while}}]]
     [[TPUReplicate/_compile/_17449455679645184466/_2]]
     [[cluster_tpu_function/control_after/_1/_111]]
  (7) INVALID_ARGUMENT: {{function_node __inference_tpu_function_230}} Input and output shapes of loop body do not match: (s32[], s32[], s32[], s32[1], s32[], /*index=5*/s32[]) vs. (s32[], s32[], s32[], s32[2], s32[], /*index=5*/s32[])

     [[{{function_node __inference_f_202}}{{node while}}]]
     [[TPUReplicate/_compile/_174 ... [truncated]

如果有办法使用 TPU 设置形状不变量，请告诉我。

Answer 1

Tensorflow 官方文档似乎给出了一个不太合适的使用 shape_invariants 循环参数的示例，这促使用户将其用作一种允许他们将循环中的中间值累加到张量中的方法。

Tensorflow 实际上有一个特殊的类 tf.TensorArray 专门设计来帮助实现这种模式。

使用 tf.TensorArray 的示例可能如下所示：

import tensorflow as tf

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

@tf.function()
def f():
  v = tf.TensorArray(tf.int32, size=3)
  for i in tf.range(3):
    v = v.write(i, i)
  return v.stack()

strategy.run(f)

我的输出：

PerReplica:{
  0: <tf.Tensor: shape=(3,), dtype=int32, numpy=array([0, 1, 2], dtype=int32)>,
  1: <tf.Tensor: shape=(3,), dtype=int32, numpy=array([0, 1, 2], dtype=int32)>,
  2: <tf.Tensor: shape=(3,), dtype=int32, numpy=array([0, 1, 2], dtype=int32)>,
  3: <tf.Tensor: shape=(3,), dtype=int32, numpy=array([0, 1, 2], dtype=int32)>,
  4: <tf.Tensor: shape=(3,), dtype=int32, numpy=array([0, 1, 2], dtype=int32)>,
  5: <tf.Tensor: shape=(3,), dtype=int32, numpy=array([0, 1, 2], dtype=int32)>,
  6: <tf.Tensor: shape=(3,), dtype=int32, numpy=array([0, 1, 2], dtype=int32)>,
  7: <tf.Tensor: shape=(3,), dtype=int32, numpy=array([0, 1, 2], dtype=int32)>
}