pytorch中torch.chunk()方法

pytorch中torch.chunk()方法

用来将张量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]]]))