torch.cat 使用小节

Septillions2023-10-16 19:37

torch.cat 要求在所指定拼接维度之外的所有维度都要匹配，例如

python 复制代码

import torch
v1 = torch.tensor([[1, 2, 3], [4, 5, 6], [4, 5, 6]])  # 3*3
v2 = torch.tensor([[3, 6, 8]])  # 1*3
torch.cat([v1, v2], dim=1)

运行之后会报错

Sizes of tensors must match except in dimension 1. Expected size 3 but got size 1 for tensor number 1 in the list.

这就是因为这两个向量的第 0 个维度不相等，故无法完成 cat 操作，改成 dim=0 即可得到输出：

python 复制代码

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

这种操作实际上完成了多个 batch 间数据的合并，若想完成同个 batch 内数据的 cat，要保证第 0 个维度大小一致，即 batchsize 相等。

python 复制代码

import torch

v1 = torch.tensor([[1, 2, 3, 3], [4, 5, 6, 6]])  # 2*4
v2 = torch.tensor([[6, 6], [8, 8]])  # 1*3
torch.cat([v1, v2], dim=1)

可得到预期输出：

python 复制代码

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