torch.split
(tensor, split_size_or_sections, dim=0)Splits the tensor into chunks. Each chunk is a view of the original 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]])
利用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])
[1] https://pytorch.org/docs/stable/generated/torch.split.html