PyTorch 的 torch.unbind 函数详解与进阶应用：中英双语

中文版

PyTorch 的 torch.unbind 函数详解与进阶应用

在深度学习中，张量的维度操作是基础又重要的内容。PyTorch 提供了许多方便的工具来完成这些操作，其中之一便是 torch.unbind 。与常见的堆叠函数（如 torch.stack）相辅相成，torch.unbind 是分解张量的重要函数。本文将从基础到进阶，详细介绍 torch.unbind 的功能及其与 torch.stack 的组合应用。

---

什么是 torch.unbind？

torch.unbind 的作用是移除指定维度，并返回一个元组，其中包含移除该维度后的张量序列。换句话说，它将张量沿指定的维度"拆开"。

函数定义

torch.unbind(input, dim=0) → Tuple[Tensor]

参数说明：

• input: 要分解的输入张量。
• dim : 指定需要移除的维度，默认是 0（第一个维度）。

返回值：

返回一个元组，每个元素是一个张量，大小为原张量去掉 dim 维度后的形状。

---

基本用法

示例 1：沿第 0 维度分解

import torch

# 创建一个 3x3 的张量
x = torch.tensor([[1, 2, 3],
                  [4, 5, 6],
                  [7, 8, 9]])

# 沿第 0 维度（行）分解
result = torch.unbind(x, dim=0)
print(result)  # 输出结果

输出：

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

分解后得到的结果是一个包含 3 个张量的元组，每个张量是原张量的一行。

---

示例 2：沿第 1 维度分解

# 沿第 1 维度（列）分解
result = torch.unbind(x, dim=1)
print(result)

输出：

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

此时每个张量是原张量的一列。

---

torch.unbind 的常见应用场景

1. 将高维数据拆分为低维数据：例如，将一个批量数据（Batch）按样本分解，或将序列数据按时间步长分解。
2. 与循环配合：分解后可以逐个张量操作，例如在处理时间序列、图像分块时非常有用。
3. 与 torch.stack 联合使用：可以将分解和重新组合操作结合起来，完成张量维度的灵活操作。

---

进阶：结合 torch.stack 使用

关于stack函数的具体用法，可参考笔者的另一篇博客：深入理解 PyTorch 中的torch.stack函数：中英双语

示例 3：分解后重新堆叠

我们可以使用 torch.unbind 将张量分解成多个子张量，然后通过 torch.stack 将这些子张量重新堆叠成新的张量。

# 创建一个 3x3 的张量
x = torch.tensor([[1, 2, 3],
                  [4, 5, 6],
                  [7, 8, 9]])

# 第一步：沿第 0 维分解
unbind_result = torch.unbind(x, dim=0)
print("分解结果：", unbind_result)

# 第二步：沿新的维度重新堆叠
stack_result = torch.stack(unbind_result, dim=1)
print("重新堆叠结果：", stack_result)

输出：

分解结果：
(tensor([1, 2, 3]), tensor([4, 5, 6]), tensor([7, 8, 9]))

重新堆叠结果：
tensor([[1, 4, 7],
        [2, 5, 8],
        [3, 6, 9]])

在这里：

• 第一步使用 torch.unbind 沿第 0 维分解，将张量拆分为 3 个张量（每个张量对应原来的行）。
• 第二步使用 torch.stack 沿第 1 维重新堆叠，生成了一个新张量，其中每个原张量成为了列。

---

示例 4：动态调整维度顺序

通过结合 torch.unbind 和 torch.stack，我们可以动态调整张量的维度顺序。

# 创建一个 3x2x2 的三维张量
x = torch.tensor([[[1, 2], [3, 4]],
                  [[5, 6], [7, 8]],
                  [[9, 10], [11, 12]]])

# 第一步：沿第 0 维分解为 3 个 2x2 张量
unbind_result = torch.unbind(x, dim=0)

# 第二步：沿第 2 维重新堆叠
stack_result = torch.stack(unbind_result, dim=2)
print("最终结果：", stack_result)

输出：

最终结果：
tensor([[[ 1,  5,  9],
         [ 3,  7, 11]],

        [[ 2,  6, 10],
         [ 4,  8, 12]]])

这里，我们通过两步实现了维度的重新排列：

1. 使用 torch.unbind 沿第 0 维分解。
2. 使用 torch.stack 沿第 2 维重新组合，从而完成了维度转换。

---

每一步变化解析

参考笔者的另一篇博客：PyTorch如何通过 torch.unbind 和torch.stack动态调整张量的维度顺序

以示例 4 为例，张量的形状在每一步的变化如下：

1. 原始张量形状为 [3, 2, 2]。
2. 分解后，得到 3 个形状为 [2, 2] 的张量。
3. 堆叠时，将这些张量沿新的维度 dim=2 组合，最终形状变为 [2, 2, 3]。

通过这种分解和堆叠方式，我们可以灵活地操作张量的维度和数据布局。

---

torch.unbind 的使用注意事项

1. 分解后返回的是元组：

   • 返回值是一个不可变的元组，而不是列表。如果需要动态修改，可以将其转换为列表。

   result = list(torch.unbind(x, dim=0))

2. 维度必须存在：

   • 如果指定的 dim 超出了张量的维度范围，PyTorch 会报错。

3. 结合其他函数使用：

   • 与 torch.stack、torch.cat、torch.split 等函数搭配，可以完成更加灵活的张量操作。

---

总结

torch.unbind 是一个高效的张量分解工具，在处理高维数据、调整维度顺序时非常有用。结合 torch.stack，它可以实现从拆分到重组的完整操作链。掌握这些函数的灵活用法，可以大大提升张量操作的效率和代码的可读性。

英文版

A Detailed Guide to PyTorch torch.unbind with Advanced Applications Using torch.stack

In deep learning, manipulating tensor dimensions is both fundamental and essential. PyTorch provides powerful tools for this purpose, and one such tool is torch.unbind . It is particularly useful for breaking down tensors into smaller components, often complementing torch.stack in advanced scenarios. This blog post provides a detailed introduction to torch.unbind, its basic functionality, and how to combine it with torch.stack for advanced applications.

---

What is torch.unbind?

The torch.unbind function removes a specified dimension from a tensor and returns a tuple containing slices of the tensor along that dimension. Simply put, it splits a tensor along a specific axis.

Function Definition

torch.unbind(input, dim=0) → Tuple[Tensor]

Parameters:

• input: The tensor to be split.
• dim : The dimension to be removed (default: 0).

Return Value:

A tuple of tensors, where each tensor is a slice of the input tensor along the specified dimension.

---

Basic Usage

Example 1: Splitting Along Dimension 0

import torch

# Create a 3x3 tensor
x = torch.tensor([[1, 2, 3],
                  [4, 5, 6],
                  [7, 8, 9]])

# Split along dimension 0 (rows)
result = torch.unbind(x, dim=0)
print(result)

Output:

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

Here, the tensor is split into three tensors, each corresponding to a row of the original tensor.

---

Example 2: Splitting Along Dimension 1

# Split along dimension 1 (columns)
result = torch.unbind(x, dim=1)
print(result)

Output:

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

This time, the tensor is split into three tensors, each corresponding to a column.

---

Common Applications of torch.unbind

1. Breaking Down Batch Data: For processing samples in a batch individually.
2. Handling Sequential Data: For splitting time steps in sequence data.
3. Dynamic Tensor Operations : Combining torch.unbind with other functions like torch.stack for flexible tensor manipulation.

---

Advanced Usage: Combining torch.unbind with torch.stack

The torch.unbind function pairs perfectly with torch.stack to achieve advanced tensor operations. Let's explore this combination.

Example 3: Split and Restack

In this example, we will split a tensor into slices and then restack it along a new dimension.

# Create a 3x3 tensor
x = torch.tensor([[1, 2, 3],
                  [4, 5, 6],
                  [7, 8, 9]])

# Step 1: Unbind along dimension 0
unbind_result = torch.unbind(x, dim=0)
print("Unbind result:", unbind_result)

# Step 2: Restack along dimension 1
stack_result = torch.stack(unbind_result, dim=1)
print("Stack result:", stack_result)

Output:

Unbind result:
(tensor([1, 2, 3]), tensor([4, 5, 6]), tensor([7, 8, 9]))

Stack result:
tensor([[1, 4, 7],
        [2, 5, 8],
        [3, 6, 9]])

Process:

1. The tensor is split into rows using torch.unbind (along dimension 0).
2. These rows are then stacked into columns using torch.stack (along dimension 1).

---

Example 4: Reordering Dimensions

Using torch.unbind and torch.stack, you can dynamically reorder tensor dimensions.

# Create a 3x2x2 tensor
x = torch.tensor([[[1, 2], [3, 4]],
                  [[5, 6], [7, 8]],
                  [[9, 10], [11, 12]]])

# Step 1: Unbind along dimension 0
unbind_result = torch.unbind(x, dim=0)

# Step 2: Stack along dimension 2
stack_result = torch.stack(unbind_result, dim=2)
print("Final result:", stack_result)

Output:

Final result:
tensor([[[ 1,  5,  9],
         [ 3,  7, 11]],

        [[ 2,  6, 10],
         [ 4,  8, 12]]])

Explanation:

1. torch.unbind splits the tensor along the first dimension (size 3), resulting in three 2×2 tensors.
2. torch.stack then combines these tensors along a new dimension (dim=2), effectively reordering the dimensions.

---

Step-by-Step Shape Transformation

Using Example 4, let's track the shape transformations:

1. Input Shape : [3, 2, 2].
2. After torch.unbind(dim=0) : A tuple of 3 tensors, each of shape [2, 2].
3. After torch.stack(dim=2) : A new tensor of shape [2, 2, 3].

---

Key Considerations

1. Output is a Tuple:

   • The result of torch.unbind is a tuple, which is immutable. To modify the result, you can convert it to a list:

   result = list(torch.unbind(x, dim=0))

2. Dimension Must Exist:

   • The specified dim must be within the tensor's dimensions; otherwise, an error will occur.

3. Paired with Other Functions:

   • torch.unbind works well with functions like torch.stack, torch.cat, and torch.split for advanced manipulation.

---

Conclusion

torch.unbind is a powerful tool for splitting tensors into smaller components, and its ability to work seamlessly with other PyTorch functions like torch.stack makes it indispensable for flexible tensor manipulation. Whether you're working on sequence data, batch processing, or dynamic dimension reordering, mastering torch.unbind will significantly enhance your PyTorch skills.

Try out these examples in your projects to fully understand the potential of torch.unbind!

后记

2024年12月12日22点20分于上海，在GPT4o大模型辅助下完成。