torch.chunk(tensor, chunks, dim)

最新推荐文章于 2024-08-11 15:44:29 发布

原创最新推荐文章于 2024-08-11 15:44:29 发布 · 496 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #pytorch #人工智能

torch.chunk函数用于将张量在指定维度上分割成多个块，chunks参数定义了块的数量。例如，将一个3x3的张量按列分成3块，会得到3个1x3的张量；按行分成2块，会得到2个3x2的张量。当块数为偶数且张量维度为奇数时，最后一块可能会包含多一列的数据。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

torch.chunk(tensor, chunks, dim)

含义：在给定维度上对所给张量进行分块。

参数解释：

Tensor -- 要进行分快的张量
chunks -- 分块个数
dim -- 维度，按照此维度进行分块

如下代码：

import torch
x = torch.randn(3, 3)
x
tensor([[ 0.5149,  1.0009, -0.7242],
        [-0.9385, -0.1157, -1.6772],
        [ 1.0187, -2.3512,  1.0539]])


torch.chunk(x, 3, dim = 0)
(tensor([[ 0.5149,  1.0009, -0.7242]]),
 tensor([[-0.9385, -0.1157, -1.6772]]),
 tensor([[ 1.0187, -2.3512,  1.0539]]))


torch.chunk(x, 3, dim = 1)
(tensor([[ 0.5149],
         [-0.9385],
         [ 1.0187]]),
 tensor([[ 1.0009],
         [-0.1157],
         [-2.3512]]),
 tensor([[-0.7242],
         [-1.6772],
         [ 1.0539]]))


torch.chunk(x, 2, dim = 1)       # 注意当chunks为偶数时，而原始张量为所分维数为奇数时的变化
(tensor([[ 0.5149,  1.0009],
         [-0.9385, -0.1157],
         [ 1.0187, -2.3512]]),
 tensor([[-0.7242],
         [-1.6772],
         [ 1.0539]]))


# 另外也可以这么使用
x.chunk(2, dim = 1)
(tensor([[ 0.5149,  1.0009],
         [-0.9385, -0.1157],
         [ 1.0187, -2.3512]]),
 tensor([[-0.7242],
         [-1.6772],
         [ 1.0539]]))