python - 如何结合 tf.map_fn 和 tf.split

Question

所以我想要的东西的伪代码是:

splitted_outputs = [tf.split(output, rate, axis=0) for output in outputs]

其中，outputs 是形状为 (512, ?, 128) 的 Tensor，splitted_outputs 是 Tensors 或 3 维 Tensor 列表的列表。所以我可以迭代这样的张量tensorflow。

我尝试使用tf.map_fn:

splitted_outputs = tf.map_fn(
    lambda output: tf.split(output, rate, axis=0),
    outputs,
    dtype=list
)

但这不可能，因为 list 不合法 tf dtype。

Answer 1

您可以使用tf.unstack在 outputs 上获取“子张量”列表，然后使用 tf.split对于每一个:

splitted_outputs = [tf.split(output, rate, axis=0) for output in tf.unstack(outputs, axis=0)]

请注意tf.unstack仅当给定 axis 的大小已知时才能这样使用，否则您需要提供 num 参数。