当前位置: 代码迷 >> 综合 >> pytorch中torch.cat(),torch.chunk(),torch.split()函数的使用方法
  详细解决方案

pytorch中torch.cat(),torch.chunk(),torch.split()函数的使用方法

热度:54   发布时间:2023-12-11 22:48:01.0

一、torch.cat()函数

熟悉C字符串的同学们应该都用过strcat()函数,这个函数在C/C++程序中用于连接2个C字符串。在pytorch中,同样有这样的函数,那就是torch.cat()函数.
先上源码定义:torch.cat(tensors,dim=0,out=None)

  • 第一个参数tensors是你想要连接的若干个张量,按你所传入的顺序进行连接,注意每一个张量需要形状相同,或者更准确的说,进行行连接的张量要求列数相同,进行列连接的张量要求行数相同
  • 第二个参数dim表示维度,dim=0则表示按行连接,dim=1表示按列连接
a=torch.tensor([[1,2,3,4],[1,2,3,4]])
b=torch.tensor([[1,2,3,4,5],[1,2,3,4,5]])
print(torch.cat((a,b),1))
#输出结果为:
tensor([[1, 2, 3, 4, 1, 2, 3, 4, 5],[1, 2, 3, 4, 1, 2, 3, 4, 5]])

二、torch.chunk()函数

torch.cat()函数是把各个tensor连接起来,这里的torch.chunk()的作用是把一个tensor均匀分割成若干个小tensor
源码定义:torch.chunk(intput,chunks,dim=0)

  • 第一个参数input是你想要分割的tensor
  • 第二个参数chunks是你想均匀分割的份数,如果该tensor在你要进行分割的维度上的size不能被chunks整除,则最后一份会略小(也可能为空)
  • 第三个参数表示分割维度,dim=0按行分割,dim=1表示按列分割
  • 该函数返回由小tensor组成的list
c=torch.tensor([[1,4,7,9,11],[2,5,8,9,13]])
print(torch.chunk(c,3,1))
#输出结果为:
(tensor([[1, 4],[2, 5]]), tensor([[7, 9],[8, 9]]), tensor([[11],[13]]))

三、torch.split()函数

这个函数可以说是torch.chunk()函数的升级版本,它不仅可以按份数均匀分割,还可以按特定方案进行分割。
源码定义:torch.split(tensor,split_size_or_sections,dim=0)

  • 第一个参数是待分割张量
  • 第二个参数有两种形式。
    一种是分割份数,这就和torch.chunk()一样了。
    第二种这是分割方案,这是一个list,待分割张量将会分割为len(list)份,每一份的大小取决于list中的元素
  • 第三个参数为分割维度
section=[1,2,1,2,2]
d=torch.randn(8,4)
print(torch.split(d,section,dim=0))
#输出结果为:
(tensor([[ 0.5388, -0.8537,  0.5539,  0.7793]]), tensor([[ 0.1270,  2.6241, -0.7594,  0.4644],[ 0.8160,  0.5553,  0.1234, -1.1157]]), tensor([[-0.4433, -0.3093, -2.0134, -0.4277]]), tensor([[-0.4297,  0.2532,  0.2789, -0.3068],[ 1.4208, -0.1202,  0.9256, -1.2127]]), tensor([[ 0.3542, -0.4656,  1.2683,  0.8753],[-0.2786, -0.2180,  0.3991,  0.5658]]))

参考文献:https://pytorch.org/docs/stable/torch.html

  相关解决方案