TensorFlow中任意形状张量按最后一维索引的通用方法求解

最简单的实现方式是使用Python原生的省略号切片语法...，它可以自动匹配张量前面所有未显式指定的维度，无需手动逐维写冒号：

# 适配所有最后一维为2的张量，前面维度数量无限制
x_0 = x[..., 0]  # 输出形状等于x的形状去掉最后一维
x_1 = x[..., 1]

如果需要一次性拆分最后一维的两个元素，也可以使用TensorFlow内置的tf.unstack方法：

x_0, x_1 = tf.unstack(x, axis=-1)

若使用tf.split实现需要注意默认会保留长度为1的最后一维，可配合tf.squeeze消去：

x_0, x_1 = tf.split(x, num_or_size_splits=2, axis=-1)
x_0 = tf.squeeze(x_0, axis=-1)
x_1 = tf.squeeze(x_1, axis=-1)

效果验证 #

• 当输入张量形状为(1,2,2)时，上述方法输出的x_0、x_1形状均为(1,2)，和示例要求一致
• 当输入张量形状为(1,2,3,4,2)时，上述方法输出的x_0、x_1形状均为(1,2,3,4)，符合预期

内容的提问来源于stack exchange，提问作者kt-kbr