### PyTorch 中 Tensor 的连接与拆分方法 #### 一、Tensor 的连接操作 在 PyTorch 中，`torch.cat()` 和 `torch.stack()` 是两种常用的张量连接方式。 1. **`torch.cat()`** - 功能：沿指定维度将多个张量拼接在一起。 - 要求：除了拼接的维度外，其他所有维度大小必须一致。 - 参数解释： - `tensors`: 需要拼接的张量列表。 - `dim`: 指定用于拼接的维度。 - 示例代码如下： ```python import torch a = torch.rand(4, 32, 8) # 形状为 [4, 32, 8] b = torch.rand(5, 32, 8) # 形状为 [5, 32, 8] c = torch.cat([a, b], dim=0) # 沿第 0 维度拼接 print(c.shape) # 输出: torch.Size([9, 32, 8]) ``` 2. **`torch.stack()`** - 功能：创建一个新的维度并将输入张量堆叠在这个新维度上。 - 特点：所有输入张量形状必须完全相同。 - 示例代码如下： ```python d = torch.stack([a, b], dim=0) # 新增一个维度作为第 0 维 print(d.shape) # 输出: torch.Size([2, 4, 32, 8]) 或者 torch.Size([2, 5, 32, 8]) ``` --- #### 二、Tensor 的拆分操作 PyTorch 提供了多种方法来实现张量的拆分功能。 1. **`torch.split()`** - 功能：按指定尺寸或分割数沿着某个维度切分张量。 - 参数解释： - `tensor`: 输入张量。 - `split_size_or_sections`: 可以是一个整数值表示每部分的长度，或者是一个列表/元组定义各部分的具体长度。 - `dim`: 切分所依据的维度。 - 示例代码如下： ```python e = torch.arange(10).reshape(5, 2) # 创建一个形状为 [5, 2] 的张量 f = torch.split(e, split_size_or_sections=2, dim=0) # 将其分为三份 for i in f: print(i) # 输出: # tensor([[0, 1], # [2, 3]]) # tensor([[4, 5], # [6, 7]]) # tensor([[8, 9]]) ``` 2. **`torch.chunk()`** - 功能：类似于 `torch.split()`，但它会尝试均匀地划分张量成若干块。 - 如果无法平均分配，则最后一块可能小于其余块。 - 示例代码如下： ```python g = torch.arange(10).reshape(5, 2) # 创建一个形状为 [5, 2] 的张量 h = torch.chunk(g, chunks=3, dim=0) # 尝试将其划分为三块 for j in h: print(j) # 输出: # tensor([[0, 1], # [2, 3]]) # tensor([[4, 5], # [6, 7]]) # tensor([[8, 9]]) ``` --- #### 总结 - 使用 `torch.cat()` 实现张量的拼接时需要注意保持除拼接维度以外的所有维度一致性[^2]^。 - 对于需要新增维度的情况可以选择 `torch.stack()` 来完成任务[^3]^。 - 当涉及到张量分解时，可以根据需求选用 `torch.split()` 或者 `torch.chunk()` 函数[^1]^。