用多个未知的尺寸重塑张量

Question

我需要能够在张量的最后一轴上重塑张量：

(None, 4) --> (None, 2, 2),

在执行时可以有如下实例：

(128, 10, 4) --> (128, 10, 2, 2)
(128, 4) --> (128, 2, 2)

是否有一个简单的解决方案，或者我需要迭代第一个轴(通过排除最后一个轴(，并考虑它可以是 None 的情况？

Answer 1

你可以这样做：

my_tensor = ...
new_shape = tf.concat([tf.shape(my_tensor)[:-1], [2, 2]], axis=0)
my_tensor_reshaped = tf.reshape(my_tensor, new_shape)