tensorflow stack unstack操作

Posted 2022-03-25 bitcarmanlee

1.stack操作

先看一下tensorflow中stack方法的函数签名

@tf_export("stack")
@dispatch.add_dispatch_support
def stack(values, axis=0, name="stack"):
  """Stacks a list of rank-`R` tensors into one rank-`(R+1)` tensor.

  See also `tf.concat`, `tf.tile`, `tf.repeat`.

  Packs the list of tensors in `values` into a tensor with rank one higher than
  each tensor in `values`, by packing them along the `axis` dimension.
  Given a list of length `N` of tensors of shape `(A, B, C)`;

  if `axis == 0` then the `output` tensor will have the shape `(N, A, B, C)`.
  if `axis == 1` then the `output` tensor will have the shape `(A, N, B, C)`.
  Etc.
  .......

stack操作是将一组秩为R的tensor叠成一个秩为R+1的tensor。

def stack_operation():
    x = tf.constant([1, 4])
    y = tf.constant([2, 5])
    z = tf.constant([3, 6])
    s1 = tf.stack([x, y, z])
    print(s1)
    s2 = tf.stack([x, y, z], axis=1)
    print(s2)

代码输出为：

tf.Tensor(
[[1 4]
 [2 5]
 [3 6]], shape=(3, 2), dtype=int32)
tf.Tensor(
[[1 2 3]
 [4 5 6]], shape=(2, 3), dtype=int32)

其中，axis指定为对哪个维度进行操作。如果axis=0，改变的则是最外层的维度，以此类推。

2.unstack

理解了上面的stack操作，unstack就好理解了，就是stack的逆操作。

def unstack_operation():
    s = tf.constant([[1, 2, 3], [4, 5, 6]])
    a = tf.unstack(s, axis=0)
    b = tf.unstack(s, axis=1)
    print(a)
    print(b)

[<tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 3], dtype=int32)>, <tf.Tensor: shape=(3,), dtype=int32, numpy=array([4, 5, 6], dtype=int32)>]
[<tf.Tensor: shape=(2,), dtype=int32, numpy=array([1, 4], dtype=int32)>, <tf.Tensor: shape=(2,), dtype=int32, numpy=array([2, 5], dtype=int32)>, <tf.Tensor: shape=(2,), dtype=int32, numpy=array([3, 6], dtype=int32)>]