Tensorflow之tf.split() 与 tf.squeeze()

最新推荐文章于 2024-01-07 12:35:32 发布

原创最新推荐文章于 2024-01-07 12:35:32 发布 · 2.2k 阅读

1 ·

CC 4.0 BY-SA版权

tensorflow 专栏收录该内容

15 篇文章

订阅专栏

本文详细解析了TensorFlow中tf.split函数的使用方法，包括参数解释、常见用例及错误处理，帮助读者深入理解如何通过该函数对张量进行分割。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

split(
    value,
    num_or_size_splits,
    axis=0,
    num=None,
    name='split'
)

输入：
value：输入的tensor
num_or_size_splits: 如果是个整数n，就将输入的tensor分为n个子tensor。如果是个tensor T，就将输入的tensor分为len(T)个子tensor。
axis：默认为0，计算value.shape[axis], 一定要能被num_or_size_splits整除。

举例：

# num_or_size_splits是tensor T，len（T）为3，所以分为3个子tensor，
# axis为1，所以value.shape[1]为30，4+15+11正好为30
split0, split1, split2 = tf.split(value, [4, 15, 11], 1)

tf.shape(split0)  # [5, 4]
tf.shape(split1)  # [5, 15]
tf.shape(split2)  # [5, 11]

# num_or_size_splits是整数3， 分为3个tensor，value.shape[1]为30，能被3整除。
split0, split1, split2 = tf.split(value, num_or_size_splits=3, axis=1)
tf.shape(split0)  # [5, 10]

再举个实例：

>>> a=np.reshape(range(24),(4,2,3))
>>> a
array([[[ 0,  1,  2],
        [ 3,  4,  5]],

       [[ 6,  7,  8],
        [ 9, 10, 11]],

       [[12, 13, 14],
        [15, 16, 17]],

       [[18, 19, 20],
        [21, 22, 23]]])

>>> sess=tf.InteractiveSession()
# 将a分为两个tensor，a.shape(1)为2，可以整除，不会报错。
# 输出应该为2个shape为[4,1,3]的tensor
>>> b= tf.split(a,2,1)
>>> b
[<tf.Tensor 'split:0' shape=(4, 1, 3) dtype=int32>, <tf.Tensor 'split:1' shape=(4, 1, 3) dtype=int32>]
>>> sess.run(b)
[array([[[ 0,  1,  2]],

       [[ 6,  7,  8]],

       [[12, 13, 14]],

       [[18, 19, 20]]]), array([[[ 3,  4,  5]],

       [[ 9, 10, 11]],

       [[15, 16, 17]],

       [[21, 22, 23]]])]
>>> c= tf.split(a,2,0)
# a.shape(0)为4，被2整除，输出2个[2,2,3]的Tensor
>>> c
[<tf.Tensor 'split_1:0' shape=(2, 2, 3) dtype=int32>, <tf.Tensor 'split_1:1' shape=(2, 2, 3) dtype=int32>]
>>> sess.run(c)
[array([[[ 0,  1,  2],
        [ 3,  4,  5]],

       [[ 6,  7,  8],
        [ 9, 10, 11]]]), array([[[12, 13, 14],
        [15, 16, 17]],

       [[18, 19, 20],
        [21, 22, 23]]])]
>>>d= tf.split(a,2,2)
#  a.shape(2)为3，不被2整除，报错。
Traceback (most recent call last):
  File "D:\Anaconda2\envs\tensorflow\lib\site-packages\tensorflow\python\framework\common_shapes.py", line 671, in _call_cpp_shape_fn_impl
    input_tensors_as_shapes, status)
  File "D:\Anaconda2\envs\tensorflow\lib\contextlib.py", line 66, in __exit__
    next(self.gen)
  File "D:\Anaconda2\envs\tensorflow\lib\site-packages\tensorflow\python\framework\errors_impl.py", line 466, in raise_exception_on_not_ok_status
    pywrap_tensorflow.TF_GetCode(status))
tensorflow.python.framework.errors_impl.InvalidArgumentError: Dimension size must be evenly divisible by 2 but is 3
        Number of ways to split should evenly divide the split dimension for 'split_1' (op: 'Split') with input shapes: [], [4,2,3] and with computed input tensors: input[0] = <2>.
>>> d= tf.split(a,3,2)
# 改成3，a.shape(2)为3，整除，不报错，返回3个[4,2,1]的Tensor
>>> d
[<tf.Tensor 'split_2:0' shape=(4, 2, 1) dtype=int32>, <tf.Tensor 'split_2:1' shape=(4, 2, 1) dtype=int32>, <tf.Tensor 'split_2:2' shape=(4, 2, 1) dtype=int32>]
>>> sess.run(d)
[array([[[ 0],
        [ 3]],

       [[ 6],
        [ 9]],

       [[12],
        [15]],

       [[18],
        [21]]]), array([[[ 1],
        [ 4]],

       [[ 7],
        [10]],

       [[13],
        [16]],

       [[19],
        [22]]]), array([[[ 2],
        [ 5]],

       [[ 8],
        [11]],

       [[14],
        [17]],

       [[20],
        [23]]])]

注意：
tf.split和reshape不同，不会改变数值之间的相对顺序。只能每个维度只能变小，不增大。
取值时按照axis-1的顺序来取。

tf.squeeze
squeeze(
    input,
    axis=None,
    name=None,
    squeeze_dims=None
)

去掉维数为1的维度。
举个例子：

# 't' is a tensor of shape [1, 2, 1, 3, 1, 1]
tf.shape(tf.squeeze(t))  # [2, 3]

也可以指定去掉哪个维度：

# 't' is a tensor of shape [1, 2, 1, 3, 1, 1]
tf.shape(tf.squeeze(t, [2, 4]))  # [1, 2, 3, 1]