tf.cond() 返回形状未知的张量

Question

下面是我传递给 keras Lambda 层的函数。

tf.cond() 的输出出现问题。它返回<unknown> 的形状。输入张量 (t) 和恒权张量的形状分别为 (None,6) 和 (6,)。当我在tf.cond() 之外添加这两个时，我得到一个形状为(None,6) 的张量，这就是我需要的。但是，当从 tf.cond() 中返回相同的添加操作时，我会得到一个形状为 <unknown> 的张量。

当此操作通过tf.cond() 进行时会发生什么变化。

def class_segmentation(t):

        class_segments = tf.constant([0,0,1,1,2,2])

        a = tf.math.segment_mean(t, class_segments, name=None)

        b = tf.math.argmax(a)
  
        left_weights = tf.constant([1.0,1.0,0.0,0.0,0.0,0.0])
        middle_weights = tf.constant([0.0,0.0,1.0,1.0,0.0,0.0])
        right_weights = tf.constant([0.0,0.0,0.0,0.0,1.0,1.0])
        zero_weights = tf.constant([0.0,0.0,0.0,0.0,0.0,0.0])

        c = tf.cond(tf.math.equal(b,0), lambda: tf.math.add(t, left_weights), lambda: zero_weights)
        d = tf.cond(tf.math.equal(b,1), lambda: tf.math.add(t, middle_weights ), lambda: zero_weights)
        e = tf.cond(tf.math.equal(b,2), lambda: tf.math.add(t, right_weights), lambda: zero_weights)

        f = tf.math.add_n([c,d,e])
        print("Tensor shape: ", f.shape) # returns "Unknown"
        
        return f

Answer 1

您的代码中存在一些问题。

1. tf.math.segment_mean() 期望 class_segments 具有与输入 t 的第一个维度相同的形状。所以None 必须等于6 才能让您的代码运行。这很可能是您获得 unknown 形状的原因 - 因为您的张量的形状取决于在运行时确定的 None。您可以应用转换来运行您的代码（不确定这是否是您想要实现的目标），例如。

a = tf.math.segment_mean(tf.transpose(t), class_segments)

1. 在tf.cond() true_fn 和false_fn 中必须返回相同形状的张量。在您的情况下，true_fn 返回 (None, 6)，因为 broadcasting 和 false_fn 返回形状为 (6,) 的张量。
2. tf.cond() 中的谓词必须降为等级 0。例如，如果您要申请 b = tf.math.argmax(tf.math.segment_mean(tf.transpose(t), class_segments), 0) 那么b 的形状将是(None)，tf.cond() 中的谓词pred 将是broadcasted 到相同的形状（这会引发错误）。

如果不知道你想获得什么进一步的帮助是不可能的。