用来将张量tensor分成很多个块,即切分,可以在不同维度上切分。返回一个张量list。
如果指定轴的元素个数被chunk数除不尽,则最后一块元素数量会少。
torch.chunk(tensor,chunk数,维度)
import torch
a=torch.tensor([[[1,2],[3,4],[3,5]],
[[5,6],[7,8],[7,9]]])
b=torch.chunk(a,2,0)
print(b)
a.shape=(2,3,2)
第0维分成两组,第0维即为最外边一层[],本身就是2 组。因此
仍然按照以前的维度
(tensor([[[1, 2],
[3, 4],
[3, 5]]]),
tensor([[[5, 6],
[7, 8],
[7, 9]]]))
b=torch.chunk(a,2,1)
第1维原本有3组,要分为两组,则第1维把3个分成2组,第一组有2个,第二组剩1 个元素
(tensor([[[1, 2], [3, 4], [3, 5]]]),
tensor([[[5, 6], [7, 8], [7, 9]]]))
import torch
a=torch.tensor([
[[1,2],[3,4],[3,5],[1,1],[5,6],[7,8],[7,9],[2,2],[9,9]]])
print(a.shape)
b=torch.chunk(a,3,1)
print(b)
9个元素分为3组,按顺序切开
输出
torch.Size([1, 9, 2])
(tensor([[[1, 2],
[3, 4],
[3, 5]]]), tensor([[[1, 1],
[5, 6],
[7, 8]]]), tensor([[[7, 9],
[2, 2],
[9, 9]]]))