如何合并尺寸不同的张量以生成指定结构的目标张量？

现有两个列数不同的PyTorch张量：

import torch

a = torch.tensor([
    [101,  103],
    [101,  1045]
])

b = torch.tensor([
    [101,  777, 227],
    [101,  888, 228]
])

想要得到如下形式的张量c：

c = torch.tensor([
    [101,  103, 0],
    [101,  1045, 0],
    [101,  777, 227],
    [101,  888, 228]
])

直接使用torch.cat((a, b), dim=0)会失败，原因是a是2行2列，b是2行3列，PyTorch要求拼接的张量除了拼接维度外，其他维度的尺寸必须一致，这里列数不匹配，所以无法直接拼接。

正确实现方法 #

先把列数少的张量a补零到和b相同的列数，再进行拼接：

方法1：用torch.nn.functional.pad补零（推荐）

# 计算需要补充的列数
pad_cols = b.size(1) - a.size(1)
# 对a的右侧补pad_cols个0
a_padded = torch.nn.functional.pad(a, (0, pad_cols), mode='constant', value=0)
# 按行拼接补零后的a和b
c = torch.cat((a_padded, b), dim=0)

方法2：手动构造零张量拼接

# 创建和a行数相同、列数为需要补充数量的零张量
zeros = torch.zeros(a.size(0), b.size(1)-a.size(1), dtype=a.dtype)
# 把a和零张量按列拼接得到补零后的a
a_padded = torch.cat((a, zeros), dim=1)
# 再按行拼接得到c
c = torch.cat((a_padded, b), dim=0)

两种方法运行后得到的c都符合预期：

tensor([[ 101,  103,    0],
        [ 101, 1045,    0],
        [ 101,  777,  227],
        [ 101,  888,  228]])

内容的提问来源于stack exchange，提问作者Brana