当我附加标签 tf.function 时，Tensorflow while_loop 抛出错误

Question

@tf.function
def manu_data(dats):
    print("Manufacturing train data...")

    dats = tf.transpose(dats, [1, 0, 2, 3])
    dats, train_adata = dats[0], dats[1]
    train_idata = tf.zeros([0, 15, 15], tf.int32)

    def body(i, train_idata, dats):
        train_idata = tf.reshape( \
                    tf.concat( \
                        [ \
                            tf.slice(train_idata, [0, 0, 0], [i, 15, 15]), \
                            tf.reshape(tf.cast(tf.divide(tf.add(tf.square(dats[i]), dats[i]), 2), tf.int32), [-1, 15, 15]), \
                            tf.reshape(tf.cast(tf.divide(tf.subtract(tf.square(dats[i]), dats[i]), 2), tf.int32), [-1, 15, 15]), \
                            tf.reshape(tf.cast(tf.subtract(1, tf.square(dats[i])), tf.int32), [-1, 15, 15]), \
                            tf.zeros([dats.shape[0] - i - 1, 15, 15], tf.int32) \
                        ], 0), \
                    [-1, 15, 15])
        if tf.equal(tf.math.mod(tf.add(i, 1), 1000), 0): tf.print(i)
        return [tf.add(i, 1), train_idata, dats]
    i = tf.constant(1)
    tf.while_loop( \
        lambda i, train_idata, dats: tf.less(i, dats.shape[0]), \
        body, [i, train_idata, dats], parallel_iterations = 16)
    train_idata = tf.reshape(train_idata, [-1, 3, 15, 15])
    return (train_idata, train_adata)

当我删除代码顶部的“@tf.function”时，我的代码运行良好。 但除此之外，我得到了这个消息：

ValueError: Input tensor 'zeros:0' enters the loop with shape (0, 15, 15), but has shape (None, 15, 15) after one iteration.

我认为这意味着张量 'train_idata' 在使用 tf.while_loop 时变成了形状 [None,15, 15] 的张量。 谁能注意到问题是什么？？

编辑：感谢 Andrea，我像这样编辑了我的代码：

tf.while_loop( \
    lambda i, train_idata, dats: tf.less(i, dats.shape[0]), \
    body, [i, train_idata, dats], parallel_iterations = 16, \
    shape_invariants = [i.shape, train_idata.shape, tf.TensorShape([None, 15, 15])])

我得到了另一个错误：

ValueError: Tried to convert 'y' to a tensor and failed. Error: None values not supported.

这意味着，我不能从具有 None 值的数组中声明 tf.TensorShape。 我应该怎么办？

你能告诉我更多细节吗？ 我是 tensorflow 新手:(

Answer 1

当使用带有 Autograph (tf.function) 的 tf.while_loop 时，所有循环变量都必须在迭代中保持形状。 在您的代码中， train_idata 以形状 (0,15,15) 开头，但在tf.concat调用期间，TF 无法推断第 0 维，因此它使用 None ，这意味着（“在运行时推断”）。 如果你想在迭代中改变形状，你应该使用shape_invariant参数。

您可以在此处阅读有关 Autograph 限制的更多信息。

Answer 2

正如Andrea Angeli 所提到的，您需要为循环提供shape_invariants参数。 除此之外，还可以在您的代码中进行一些简化。 新行末尾不需要任何\ （如果您在括号、方括号或大括号内，则可以在 Python 中换行）。 大多数 TensorFlow 操作可以使用常规数学运算符以更简单的方式表示。 还有一些其他细节。 以下是它的工作原理：

import tensorflow as tf

@tf.function
def manu_data(dats):
    # Use tf.print so it gets printed on every call,
    # not just the first one
    tf.print("Manufacturing train data...")

    dats = tf.transpose(dats, [1, 0, 2, 3])
    dats, train_adata = dats[0], dats[1]
    train_idata = tf.zeros([0, 15, 15], tf.int32)

    def body(i, train_idata, dats):
        train_idata = tf.reshape(
            tf.concat([
                train_idata[:i, :15, :15],
                tf.reshape(tf.cast((tf.square(dats[i]) + dats[i]) / 2, tf.int32), [-1, 15, 15]),
                tf.reshape(tf.cast(tf.square(dats[i]) - dats[i] / 2, tf.int32), [-1, 15, 15]),
                tf.reshape(tf.cast(1 - tf.square(dats[i]), tf.int32), [-1, 15, 15]),
                tf.zeros([tf.shape(dats)[0] - i - 1, 15, 15], tf.int32)], axis=0),
            [-1, 15, 15])
        # Change printing to be "graph-friendly"
        maybe_print = tf.cond(tf.equal(i + 1 % 1000, 0),
            lambda: tf.print(i),
            lambda: tf.group([]))  # No-op
        with tf.control_dependencies([maybe_print]):
            return [i + 1, train_idata, dats]
    i = tf.constant(1)
    _, train_idata, _ = tf.while_loop(
        # Better use tf.shape to support dynamic shape
        lambda i, train_idata, dats: i < tf.shape(dats)[0],
        body,
        [i, train_idata, dats],
        # Use shape invariants to support variable shapes
        shape_invariants=[tf.TensorShape(()), tf.TensorShape((None, 15, 15)), dats.shape],
        parallel_iterations=16)
    train_idata = tf.reshape(train_idata, [-1, 3, 15, 15])
    return (train_idata, train_adata)

# Test
out = manu_data(tf.zeros((100, 2, 15, 15), tf.float32))

无论如何，您都没有收集tf.while_loop的结果，所以我认为您可能误解了它的目的。 查看 function 的文档和示例，您应该将其产生的输出作为结果 - 可能您需要类似_, train_idata, _ = tf.while_loop(...)的内容。 此外，通常不建议在这样的循环中向张量添加行，在大多数情况下，有更好的方法来做同样的事情。

Answer 3

我不确定原始代码打算计算什么，但在这里它被简化了一点，以利用定义的运算符，从而减少由于 function 调用造成的混乱。 考虑到train_idata的初始值，我认为您还希望i从 0 开始，而不是 1：

@tf.function
def manu_data(dats):
    tf.print("Manufacturing train data...")

    dats = tf.transpose(dats, [1, 0, 2, 3])
    dats, train_adata = dats[0], dats[1]
    train_idata = tf.zeros([0, 15, 15], tf.int32)

    i = 0
    while i < dats.shape[0]:
      tf.autograph.experimental.set_loop_options(
          parallel_iterations=16
      )
      train_idata = tf.reshape(
          tf.concat(
              [
               train_idata[0:i, 0:15, 0:15],
               tf.reshape(tf.cast(tf.square(dats[i]) + dats[i] / 2, tf.int32), [-1, 15, 15]),
               tf.reshape(tf.cast(tf.square(dats[i]) - dats[i] / 2, tf.int32), [-1, 15, 15]),
               tf.reshape(tf.cast(1 - tf.square(dats[i]), tf.int32), [-1, 15, 15]),
               tf.zeros([dats.shape[0] - i - 1, 15, 15], tf.int32)
              ],
              0),
              [-1, 15, 15])

      if i + 1 % 1000 == 0:
        tf.print(i)

      i += 1

    train_idata = tf.reshape(train_idata, [-1, 15, 15])
    return train_idata, train_adata

manu_data(tf.ones([30, 2, 15, 15], dtype=tf.float32))

这在我的测试中运行没有错误，但我不确定它是否具有您期望的输入/输出。