PyTorch中如何扩展张量维度得到指定结构的张量

• 原张量A是shape为(2, 3)的二维张量
• 目标张量B是shape为(2, 2, 3)的三维张量，两个最外层切片的内容和A完全一致

以下是三种可直接运行的实现方案，可根据场景选择：

方法1：unsqueeze + repeat（生成独立内存副本，最常用） #

repeat会真实复制张量数据，生成的新张量和原张量完全独立，修改值互不影响：

import torch
A = torch.tensor([[1, 2, 3], [4, 5, 6]])
# 先在第0维新增长度为1的维度，再在该维度复制2次，其余维度保持原长度
B = A.unsqueeze(0).repeat(2, 1, 1)

参数说明：repeat()传入的参数对应每个维度的复制次数，传入1代表对应维度不做复制。

方法2：unsqueeze + expand（内存占用低，返回广播视图） #

expand不会复制真实数据，仅返回原张量的广播视图，内存开销极小，适合不需要修改张量值的只读场景：

# 先升维，再把第0维长度扩展到2，传入-1代表对应维度保持原长度不变
B = A.unsqueeze(0).expand(2, -1, -1)

注意：该方法返回的张量和原A共享底层内存，修改B的元素会同步改动A的值，如果需要独立副本可以在末尾加.clone()：B = A.unsqueeze(0).expand(2, -1, -1).clone()

方法3：torch.stack堆叠（写法直观，新手不易出错） #

不需要手动计算维度复制参数，直接把两份原张量在新维度上堆叠即可，生成的是独立内存的新张量：

# dim=0代表在最外层新增维度，堆叠两份A
B = torch.stack([A, A], dim=0)

结果验证 #

运行以下代码可以确认输出完全符合预期：

print(B)
print(B.shape)

输出内容：

tensor([[[1, 2, 3],
         [4, 5, 6]],

        [[1, 2, 3],
         [4, 5, 6]]])
torch.Size([2, 2, 3])

内容的提问来源于stack exchange，提问作者Pengfei.C