torch.split()

torch.split(tensor, split_size_or_sections, dim=0)

Splits the tensor into chunks. Each chunk is a view of the original tensor.

• tensor (Tensor) – tensor to split. 输入

• split_size_or_sections (int) or (list(int)) – size of a single chunk or list of sizes for each chunk

  split_size_or_sections可以是整数也可以是list，int时每个chunk的大小；list里面每个元素对应一个chunk的大小。

• dim (int) – dimension along which to split the tensor. 切分维度

import torch

a = torch.arange(10).reshape(5,2)
print(a.shape)
aa = torch.split(a, 2)
print(aa)
b,c,d = torch.split(a, 2)
print(b.shape,c.shape,d.shape)
print('*' * 30)
e,f = torch.split(a, [1,4]) # 沿着dim=0的方向，切分为1,4大小的chunk
print(e.shape, f.shape)
print(e,f)

Result:

torch.Size([5, 2])
(tensor([[0, 1],
        [2, 3]]), tensor([[4, 5],
        [6, 7]]), tensor([[8, 9]]))
torch.Size([2, 2]) torch.Size([2, 2]) torch.Size([1, 2])
******************************
torch.Size([1, 2]) torch.Size([4, 2])
tensor([[0, 1]]) tensor([[2, 3],
        [4, 5],
        [6, 7],
        [8, 9]])

Example

利用torch.split()来对图片按通道切片

# SHM 
trimap = torch.randn(1,3,512,512)
trimap_softmax = F.softmax(trimap, dim=1) ###
print(trimap_softmax.shape)
# paper: bs, fs, us
bg, fg, unsure = torch.split(trimap_softmax, 1, dim=1)
print(bg.shape, fg.shape, unsure.shape)

Result:

torch.Size([1, 3, 512, 512])
torch.Size([1, 1, 512, 512]) torch.Size([1, 1, 512, 512]) torch.Size([1, 1, 512, 512])

Reference:

[1] https://pytorch.org/docs/stable/generated/torch.split.html