尝试在 tensorflow2 中连接批量大小不一致的张量答案

【问题标题】：Try to concatenate tensors of inconsistent batch size in tensorflow2尝试在 tensorflow2 中连接批量大小不一致的张量
【发布时间】：2021-11-18 17:20:12
【问题描述】：

我尝试在一个用@tf.function 修饰的函数中连接 2 个不同批量大小的张量。我尝试了两种方法，第一种方法如下：

import tensorflow as tf

@tf.function  # indispensable
def fun1(tensors, indices):
    results = []

    for i in tf.range(2):  # batch size = 2
        pos = tf.where(indices==i)
        emb = tf.gather_nd(tensors, pos)
        # do something to emb, but do nothing here for simplicity.
        results += [emb]

    results = tf.concat(results, axis=0)
    return results

tensors = tf.random.uniform((5, 2))
fun1(tensors, indices=[0, 0, 1, 1, 1])

但它会引发如下错误：

TypeError: 'results' does not have the same nested structure after one iteration.
The two structures don't have the same nested structure.
First structure: type=list str=[]
Second structure: type=list str=[<tf.Tensor 'while/GatherNd:0' shape=(None, 2) dtype=float32>]

More specifically: The two structures don't have the same number of elements. First structure: type=list str=[]. Second structure: type=list str=[<tf.Tensor 'while/GatherNd:0' shape=(None, 2) dtype=float32>]
Entire first structure:
[]
Entire second structure:
[.]

于是我尝试了第二种方法：

import tensorflow as tf

@tf.function  # indispensable
def fun2(tensors, indices):
    results = tf.reshape(tf.constant([], dtype=tf.float32), (0, 2))  # make empty tensors

    for i in tf.range(2):  # batch size = 2
        pos = tf.where(indices==i)
        emb = tf.gather_nd(tensors, pos)
        # do something to emb, but do nothing here for simplicity
        results = tf.concat([results, emb], axis=0)

    return results

tensors = tf.random.uniform((5, 2))
fun2(tensors, indices=[0, 0, 1, 1, 1])

但它会引发错误：

ValueError: 'results' has shape (0, 2) before the loop, but shape (None, 2) after one iteration. Use tf.autograph.experimental.set_loop_options to set shape invariants.

我应该如何解决这些问题？谢谢

【问题讨论】：

标签： tensorflow2.x

【解决方案1】：

我发现我可以通过在第二种方法中添加一行代码来实现它，如下所示：

@tf.function
def fun2(tensors, indices):
    results = tf.reshape(tf.constant([], dtype=tf.float32), (0, 2))  # make empty tensors

    for i in tf.range(2):  # batch size = 2
        tf.autograph.experimental.set_loop_options(shape_invariants=[(results, tf.TensorShape([None, 2]))])

        pos = tf.where(indices==i)
        emb = tf.gather_nd(tensors, pos)
        # do something to emb, but do nothing here for simplicity
        results = tf.concat([results, emb], axis=0)

    return results

【讨论】：