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
本文参与 腾讯云自媒体同步曝光计划，分享自作者个人站点/博客。
原始发表：2021/01/08 ，如有侵权请联系 cloudcommunity@tencent.com 删除
评论
登录后参与评论
推荐阅读
目录
腾讯云开发者
