pytorch中torch.stack()的功能

首先看一下官方文档对torch.stack()函数功能的描述：

       

 将一系列tensor沿着dim维度方向串联起来，所有的tensor必须具备相同的大小，其中维度的大小必须是在0到需要被串联的tensor的维度大小之间，如被串联的tensor是一个二维平面，则dim不能大于2。

简单来说：就是把多个2维张量串联成3维张量，把多个3维张量串联为4维张量，以此类推就是在增加新的维度进行堆叠。

举个例子：

将3个2维张量进行torch.stack()操作。首先准备3个2维的3行4列数据如下图：

  

数据形状如下：

 

a,b 在0维上串联后得到的新tensor大小为（2，3，4）；

a,b 在1维上串联后得到的新tensor大小为（3，2，4）。