企业官网建设流程全解析

【Bug已解决】How to use stack() in PyTorch? 解决方案

问题描述

在 PyTorch 中，torch.stack()是一个常用的张量拼接函数，用于将多个张量沿新的维度堆叠在一起。然而，许多开发者对stack()和cat()的区别不清楚，在使用时遇到维度不匹配、堆叠方向错误等问题。

常见问题包括：

• stack()和cat()有什么区别？
• 如何沿指定维度堆叠张量？
• 为什么stack()报维度不匹配错误？
• 如何堆叠不同形状的张量？
• 如何在训练循环中收集 batch 的输出？
• stack()对内存有什么影响？

torch.stack()沿一个新的维度堆叠张量，而torch.cat()沿已有维度拼接张量。理解这个区别是正确使用这两个函数的关键。

错误复现

场景一：stack 和 cat 混淆

import torch a = torch.tensor([1, 2, 3]) b = torch.tensor([4, 5, 6]) # stack: 沿新维度堆叠 stacked = torch.stack([a, b]) print(f"stack 结果: {stacked}") # tensor([[1, 2, 3], # [4, 5, 6]]) 形状 (2, 3) # cat: 沿已有维度拼接 catted = torch.cat([a, b]) print(f"cat 结果: {catted}") # tensor([1, 2, 3, 4, 5, 6]) 形状 (6,)

场景二：形状不匹配

a = torch.tensor([1, 2, 3]) # (3,) b = torch.tensor([4, 5]) # (2,) try: result = torch.stack([a, b]) except RuntimeError as e: print(f"错误: {e}") # stack expects each tensor to be equal size, but got [3] at entry 0 and [2] at entry 1

场景三：维度理解错误

a = torch.randn(3, 4) b = torch.randn(3, 4) # 沿 dim=0 堆叠 result0 = torch.stack([a, b], dim=0) print(f"dim=0 形状: {result0.shape}") # (2, 3, 4) # 沿 dim=1 堆叠 result1 = torch.stack([a, b], dim=1) print(f"dim=1 形状: {result1.shape}") # (3, 2, 4) # 沿 dim=2 堆叠 result2 = torch.stack([a, b], dim=2) print(f"dim=2 形状: {result2.shape}") # (3, 4, 2)

场景四：训练循环中收集输出

# 错误: 在循环中不断 stack，效率低 outputs = None for batch in dataloader: output = model(batch) if outputs is None: outputs = output else: outputs = torch.stack([outputs, output]) # 每次都创建新张量 # 正确: 先收集到列表，最后一次性 stack outputs_list = [] for batch in dataloader: output = model(batch) outputs_list.append(output) outputs = torch.stack(outputs_list)

根因分析

1. stack vs cat 的本质区别

• torch.stack(tensors, dim)：沿一个新的维度堆叠，所有张量形状必须完全相同。输出比输入多一个维度。
• torch.cat(tensors, dim)：沿已有维度拼接，除拼接维度外其他维度必须相同。输出与输入维度数相同。

a = torch.randn(3, 4) # 形状 (3, 4) b = torch.randn(3, 4) # 形状 (3, 4) # stack: 输出多一个维度 torch.stack([a, b], dim=0).shape # (2, 3, 4) - 新增维度 0 torch.stack([a, b], dim=1).shape # (3, 2, 4) - 新增维度 1 torch.stack([a, b], dim=2).shape # (3, 4, 2) - 新增维度 2 # cat: 维度数不变 torch.cat([a, b], dim=0).shape # (6, 4) - dim 0 变长 torch.cat([a, b], dim=1).shape # (3, 8) - dim 1 变长

2. stack 的工作原理

torch.stack([a, b], dim)等价于：

1. 先对每个张量unsqueeze添加新维度
2. 然后沿新维度cat

a = torch.randn(3, 4) b = torch.randn(3, 4) # stack([a, b], dim=0) 等价于: result = torch.cat([a.unsqueeze(0), b.unsqueeze(0)], dim=0) # (2, 3, 4) # stack([a, b], dim=1) 等价于: result = torch.cat([a.unsqueeze(1), b.unsqueeze(1)], dim=1) # (3, 2, 4)

3. 形状要求

stack要求所有输入张量形状完全相同：

# 正确: 形状相同 a = torch.randn(3, 4) b = torch.randn(3, 4) torch.stack([a, b]) # OK # 错误: 形状不同 a = torch.randn(3, 4) b = torch.randn(3, 5) torch.stack([a, b]) # RuntimeError

4. dim 参数的含义

dim指定新维度插入的位置。对于 N 维输入张量，dim的范围是[-(N+1), N]：

a = torch.randn(3, 4) # 2 维 # dim 范围: [-3, 2] torch.stack([a, a], dim=-3) # 等价于 dim=0 torch.stack([a, a], dim=-2) # 等价于 dim=1 torch.stack([a, a], dim=-1) # 等价于 dim=2

解决方案

方案一：基本堆叠操作

import torch # 1D 张量堆叠 a = torch.tensor([1, 2, 3]) b = torch.tensor([4, 5, 6]) c = torch.tensor([7, 8, 9]) # 默认 dim=0 result = torch.stack([a, b, c]) print(f"默认堆叠: {result.shape}") # (3, 3) print(result) # tensor([[1, 2, 3], # [4, 5, 6], # [7, 8, 9]]) # dim=1 result = torch.stack([a, b, c], dim=1) print(f"dim=1: {result.shape}") # (3, 3) print(result) # tensor([[1, 4, 7], # [2, 5, 8], # [3, 6, 9]]) # 2D 张量堆叠 x = torch.randn(3, 4) y = torch.randn(3, 4) result0 = torch.stack([x, y], dim=0) print(f"dim=0: {result0.shape}") # (2, 3, 4) result1 = torch.stack([x, y], dim=1) print(f"dim=1: {result1.shape}") # (3, 2, 4) result2 = torch.stack([x, y], dim=2) print(f"dim=2: {result2.shape}") # (3, 4, 2)

方案二：训练循环中收集输出

import torch import torch.nn as nn model = nn.Linear(10, 5) # 模拟多个 batch batches = [torch.randn(4, 10) for _ in range(5)] # 方法 A: 收集到列表，最后 stack outputs = [] for batch in batches: output = model(batch) outputs.append(output) all_outputs = torch.stack(outputs, dim=0) print(f"所有输出形状: {all_outputs.shape}") # (5, 4, 5) # dim 0: batch 索引, dim 1: 样本, dim 2: 特征 # 方法 B: 使用列表推导式 all_outputs = torch.stack([model(batch) for batch in batches]) print(f"列表推导式: {all_outputs.shape}") # (5, 4, 5)

方案三：堆叠不同维度的张量

# 堆叠 3D 张量 a = torch.randn(2, 3, 4) b = torch.randn(2, 3, 4) # 沿不同维度堆叠 for dim in range(4): result = torch.stack([a, b], dim=dim) print(f"dim={dim}: {result.shape}") # dim=0: (2, 2, 3, 4) # dim=1: (2, 2, 3, 4) # dim=2: (2, 3, 2, 4) # dim=3: (2, 3, 4, 2)

方案四：处理形状不完全相同的张量

# 如果张量形状不完全相同，需要先对齐 a = torch.randn(3, 4) b = torch.randn(3, 5) # 第二维不同 # 方法 A: 使用 cat 而非 stack result = torch.cat([a, b], dim=1) print(f"cat dim=1: {result.shape}") # (3, 9) # 方法 B: 填充到相同形状再 stack max_dim1 = max(a.shape[1], b.shape[1]) a_padded = torch.zeros(3, max_dim1) a_padded[:, :a.shape[1]] = a b_padded = torch.zeros(3, max_dim1) b_padded[:, :b.shape[1]] = b result = torch.stack([a_padded, b_padded]) print(f"填充后 stack: {result.shape}") # (2, 3, 5)

方案五：stack 与 unstack（split）

# stack 的逆操作 x = torch.randn(3, 4) y = torch.randn(3, 4) z = torch.randn(3, 4) # stack stacked = torch.stack([x, y, z], dim=0) # (3, 3, 4) print(f"stacked: {stacked.shape}") # unstack (使用 unbind) unstacked = torch.unbind(stacked, dim=0) print(f"unstacked 数量: {len(unstacked)}") # 3 print(f"每个形状: {unstacked[0].shape}") # (3, 4) # 验证 for i, t in enumerate(unstacked): print(f"还原 {i}: {torch.allclose(t, [x, y, z][i])}") # True

完整修复代码

以下是一个完整的 stack 操作工具模块：

![配图](https://i-blog.csdnimg.cn/img_convert/4e74e3a5a613e258abd1e1f8a33e1c43.png) """ PyTorch stack 操作工具模块 """ import torch from typing import List, Optional class StackUtils: """张量堆叠工具类""" @staticmethod def safe_stack(tensors: List[torch.Tensor], dim: int = 0, pad_value: float = 0.0) -> torch.Tensor: """ 安全堆叠，自动处理形状不一致的情况 Args: tensors: 张量列表 dim: 堆叠维度 pad_value: 填充值 Returns: 堆叠后的张量 """ if not tensors: raise ValueError("张量列表为空") # 检查形状是否一致 shapes = [t.shape for t in tensors] if all(s == shapes[0] for s in shapes): return torch.stack(tensors, dim=dim) # 形状不一致，需要填充 print(f"警告: 张量形状不一致 {shapes}，进行填充") # 计算每维度的最大值 ndim = tensors[0].ndim max_shape = [] for d in range(ndim): max_shape.append(max(t.shape[d] for t in tensors)) # 填充每个张量 padded = [] for t in tensors: pad_dims = [] for d in range(ndim - 1, -1, -1): diff = max_shape[d] - t.shape[d] pad_dims.extend([0, diff]) padded_t = torch.nn.functional.pad(t, pad_dims, value=pad_value) padded.append(padded_t) return torch.stack(padded, dim=dim) @staticmethod def stack_with_info(tensors: List[torch.Tensor], dim: int = 0, names: Optional[List[str]] = None) -> torch.Tensor: """ 堆叠张量并打印详细信息 Args: tensors: 张量列表 dim: 堆叠维度 names: 张量名称列表（用于调试） """ print(f"--- Stack 操作 ---") print(f"堆叠维度: dim={dim}") print(f"张量数量: {len(tensors)}") for i, t in enumerate(tensors): name = names[i] if names else f"tensor_{i}" print(f" {name}: shape={t.shape}, dtype={t.dtype}") # 检查形状一致性 shapes = [t.shape for t in tensors] if not all(s == shapes[0] for s in shapes): print(f" [错误] 形状不一致: {shapes}") raise RuntimeError("所有张量形状必须一致") result = torch.stack(tensors, dim=dim) print(f"结果: shape={result.shape}, dtype={result.dtype}") print() return result @staticmethod def collect_and_stack(items: List[torch.Tensor], dim: int = 0) -> torch.Tensor: """ 收集列表中的张量并堆叠 常用于训练循环中收集 batch 输出 """ if not items: raise ValueError("列表为空") return torch.stack(items, dim=dim) @staticmethod def unstack(tensor: torch.Tensor, dim: int = 0) -> List[torch.Tensor]: """ 拆分堆叠的张量（stack 的逆操作） """ return list(torch.unbind(tensor, dim=dim)) @staticmethod def restack(tensor: torch.Tensor, from_dim: int, to_dim: int) -> torch.Tensor: """ 将张量从一个堆叠维度移动到另一个维度 Args: tensor: 输入张量 from_dim: 当前的堆叠维度 to_dim: 目标堆叠维度 Returns: 重新堆叠的张量 """ # 先 unstack parts = StackUtils.unstack(tensor, dim=from_dim) # 再 stack 到新维度 return torch.stack(parts, dim=to_dim) # ============================================================ # 完整示例 # ============================================================ def demo_basic_stack(): """基本 stack 操作""" print("=" * 60) print("基本 stack 操作") print("=" * 60) # 1D 堆叠 a = torch.tensor([1, 2, 3]) b = torch.tensor([4, 5, 6]) print("\n--- 1D 堆叠 ---") print(f"dim=0: {torch.stack([a, b], dim=0)}") print(f"dim=1: {torch.stack([a, b], dim=1)}") # 2D 堆叠 x = torch.randn(2, 3) y = torch.randn(2, 3) print("\n--- 2D 堆叠 ---") for dim in range(3): result = torch.stack([x, y], dim=dim) print(f"dim={dim}: {result.shape}") # 3D 堆叠 a3d = torch.randn(2, 3, 4) b3d = torch.randn(2, 3, 4) print("\n--- 3D 堆叠 ---") for dim in range(4): result = torch.stack([a3d, b3d], dim=dim) print(f"dim={dim}: {result.shape}") def demo_stack_vs_cat(): """stack vs cat 对比""" print("\n" + "=" * 60) print("stack vs cat 对比") print("=" * 60) a = torch.randn(3, 4) b = torch.randn(3, 4) print(f"\n输入: a{a.shape}, b{b.shape}") # stack 操作 print("\n--- torch.stack ---") for dim in range(3): result = torch.stack([a, b], dim=dim) print(f" stack dim={dim}: {result.shape}") # cat 操作 print("\n--- torch.cat ---") for dim in range(2): result = torch.cat([a, b], dim=dim) print(f" cat dim={dim}: {result.shape}") def demo_training_loop(): """训练循环中的 stack 使用""" print("\n" + "=" * 60) print("训练循环中的 stack 使用") print("=" * 60) model = torch.nn.Sequential( torch.nn.Linear(10, 32), torch.nn.ReLU(), torch.nn.Linear(32, 5) ) # 模拟 5 个 batch batches = [torch.randn(4, 10) for _ in range(5)] labels = [torch.randint(0, 5, (4,)) for _ in range(5)] # 收集所有 batch 的输出 all_outputs = [] all_losses = [] criterion = torch.nn.CrossEntropyLoss() for i, (batch, label) in enumerate(zip(batches, labels)): output = model(batch) loss = criterion(output, label) all_outputs.append(output) all_losses.append(loss) # 堆叠 stacked_outputs = torch.stack(all_outputs, dim=0) stacked_losses = torch.stack(all_losses, dim=0) print(f"batch 数量: {len(all_outputs)}") print(f"堆叠输出形状: {stacked_outputs.shape}") # (5, 4, 5) print(f"堆叠损失形状: {stacked_losses.shape}") # (5,) print(f"平均损失: {stacked_losses.mean().item():.4f}") # 也可以沿其他维度堆叠 # 如果想按样本堆叠 stacked_by_sample = torch.stack(all_outputs, dim=1) print(f"按样本堆叠: {stacked_by_sample.shape}") # (4, 5, 5) def demo_safe_stack(): """安全堆叠（处理形状不一致）""" print("\n" + "=" * 60) print("安全堆叠（处理形状不一致）") print("=" * 60) # 形状一致 - 正常堆叠 a = torch.randn(3, 4) b = torch.randn(3, 4) result = StackUtils.safe_stack([a, b]) print(f"形状一致: {result.shape}") # (2, 3, 4) # 形状不一致 - 自动填充 a = torch.randn(3, 4) b = torch.randn(3, 6) result = StackUtils.safe_stack([a, b], pad_value=-1) print(f"形状不一致（填充后）: {result.shape}") # (2, 3, 6) def demo_unstack_restack(): """unstack 和 restack""" print("\n" + "=" * 60) print("unstack 和 restack") print("=" * 60) # 创建堆叠张量 a = torch.randn(3, 4) b = torch.randn(3, 4) c = torch.randn(3, 4) stacked = torch.stack([a, b, c], dim=0) # (3, 3, 4) print(f"原始堆叠: {stacked.shape}") # unstack parts = StackUtils.unstack(stacked, dim=0) print(f"unstack 后: {len(parts)} 个张量, 每个 {parts[0].shape}") # 验证还原 print(f"还原正确: {torch.allclose(parts[0], a)}") # restack 到不同维度 restacked = StackUtils.restack(stacked, from_dim=0, to_dim=2) print(f"restack dim=0->2: {restacked.shape}") # (3, 4, 3) # 验证 parts2 = StackUtils.unstack(restacked, dim=2) print(f"restack 后还原: {torch.allclose(parts2[0], a)}") def demo_practical_use(): """实际应用场景""" print("\n" + "=" * 60) print("实际应用场景") print("=" * 60) # 场景: 多模型集成 print("\n--- 多模型集成 ---") models = [torch.nn.Linear(10, 5) for _ in range(3)] x = torch.randn(4, 10) # 收集所有模型的输出 outputs = [model(x) for model in models] stacked = torch.stack(outputs, dim=0) # (3, 4, 5) print(f"3 个模型输出堆叠: {stacked.shape}") # 集成预测（平均） ensemble_pred = stacked.mean(dim=0) # (4, 5) print(f"集成预测: {ensemble_pred.shape}") # 场景: 多通道特征 print("\n--- 多通道特征 ---") features = [torch.randn(32, 64) for _ in range(5)] # 5 种特征 multi_channel = torch.stack(features, dim=1) # (32, 5, 64) print(f"多通道特征: {multi_channel.shape}") # 场景: 时间序列 print("\n--- 时间序列 ---") time_steps = [torch.randn(8, 10) for _ in range(6)] # 6 个时间步 sequence = torch.stack(time_steps, dim=1) # (8, 6, 10) print(f"时间序列: {sequence.shape}") if __name__ == "__main__": demo_basic_stack() demo_stack_vs_cat() demo_training_loop() demo_safe_stack() demo_unstack_restack() demo_practical_use()

常见陷阱与注意事项

1. stack 要求所有张量形状完全相同

# 正确 a = torch.randn(3, 4) b = torch.randn(3, 4) torch.stack([a, b]) # OK # 错误 a = torch.randn(3, 4) b = torch.randn(4, 3) torch.stack([a, b]) # RuntimeError

2. stack 增加一个维度，cat 不增加

a = torch.randn(3, 4) # 2D b = torch.randn(3, 4) torch.stack([a, b]).shape # (2, 3, 4) - 3D torch.cat([a, b]).shape # (6, 4) - 2D (默认 dim=0)

3. dim 参数的范围

a = torch.randn(3, 4) # 2D 张量 # dim 范围: [-3, 2] torch.stack([a, a], dim=0) # (2, 3, 4) - OK torch.stack([a, a], dim=1) # (3, 2, 4) - OK torch.stack([a, a], dim=2) # (3, 4, 2) - OK torch.stack([a, a], dim=-1) # (3, 4, 2) - 等价于 dim=2 torch.stack([a, a], dim=-3) # (2, 3, 4) - 等价于 dim=0 # torch.stack([a, a], dim=3) # IndexError - 超出范围

4. 避免在循环中反复 stack

# 错误: 每次迭代都创建新张量 result = None for item in items: if result is None: result = item.unsqueeze(0) else: result = torch.stack([result, item]) # O(n^2) 复杂度 # 正确: 先收集到列表 result_list = [] for item in items: result_list.append(item) result = torch.stack(result_list) # O(n) 复杂度

5. stack 的逆操作是 unbind

# stack 和 unbind 是互逆操作 a = torch.randn(3, 4) b = torch.randn(3, 4) stacked = torch.stack([a, b], dim=0) # (2, 3, 4) unstacked = torch.unbind(stacked, dim=0) # [tensor(3,4), tensor(3,4)] torch.allclose(unstacked[0], a) # True torch.allclose(unstacked[1], b) # True

6. 数据类型一致性

a = torch.randn(3, 4, dtype=torch.float32) b = torch.randn(3, 4, dtype=torch.float64) # stack 会自动提升类型 result = torch.stack([a, b]) print(result.dtype) # torch.float64

7. GPU 上的 stack

if torch.cuda.is_available(): a = torch.randn(3, 4, device='cuda') b = torch.randn(3, 4, device='cuda') result = torch.stack([a, b]) # 在 GPU 上 print(result.device) # cuda:0 # 不能混合 CPU 和 GPU 张量 # c = torch.randn(3, 4) # CPU # torch.stack([a, c]) # RuntimeError

总结

torch.stack()是 PyTorch 中沿新维度堆叠张量的重要操作，理解其与torch.cat()的区别是正确使用的关键。

核心要点总结：

1. stack vs cat：stack沿新维度堆叠（输出多一个维度），cat沿已有维度拼接（维度数不变）。所有 stack 的输入张量形状必须完全相同。

2. dim 参数：指定新维度插入的位置。对于 N 维输入，dim 范围是[-(N+1), N]。dim=0在最前面插入新维度。

3. stack 的等价操作：stack([a, b], dim=d)等价于先unsqueeze再cat：cat([a.unsqueeze(d), b.unsqueeze(d)], dim=d)。

4. 训练循环中使用：先收集输出到列表，最后一次性stack，避免在循环中反复 stack 导致 O(n²) 复杂度。

5. 逆操作：torch.unbind(tensor, dim)是stack的逆操作，将堆叠的张量拆分回列表。

6. 形状不一致处理：如果张量形状不完全相同，需要先填充对齐再 stack，或使用cat代替。

7. 实际应用：多模型集成（dim=0 堆叠各模型输出）、多通道特征（dim=1 堆叠不同特征）、时间序列（dim=1 堆叠时间步）。

8. 内存注意：stack 会创建新张量，对于大量大张量，注意内存消耗。可以使用分批处理。

通过掌握torch.stack()的正确用法，你可以高效地组织和处理多个张量，满足各种深度学习场景的需求。