【Python进阶】einops库的总：张量重排，重复，堆叠操作利器

【Python进阶】einops库的总：张量重排，重复，堆叠操作利器 | rearrange reduce repeat stack

深度学习中，往往需要大量的形状改变，我见到的最多的便是rearrange

他其实是属于einops这个库的，这个库的更多使用方式整理如下

文章目录

1 rearrange函数重排
2 reduce 计算平均值
3 repeat和stack

einops 是一个用于重新排列和重塑张量的库，它的目标是简化张量操作的表达和实现。 einops 提供了一种类似于 Einstein 符号约定的语法，用于描述张量操作，使得代码更具可读性和可维护性。

einops 的核心是 einops.rearrange 函数，该函数接受一个输入张量和一个描述如何重新排列张量的字符串，并返回重新排列后的张量。除了 rearrange 函数之外，einops 还提供了许多其他函数，如 einops.reduce、einops.repeat、einops.stack 等，用于执行各种张量操作。

1 rearrange函数重排

当我们需要将形状为 (batch_size, height, width, channels) 的张量重排成 (batch_size, channels, height, width) 时，可以使用 einops 的 rearrange 函数。示例如下：

import torch
from einops import rearrange# 创建一个形状为 (2, 3, 4, 5) 的张量
x = torch.randn(2, 3, 4, 5)# 使用 einops 将张量重排为 (2, 5, 3, 4)
x_rearranged = rearrange(x, 'batch height width channels -> batch channels height width')print(x.shape)             # 输出: torch.Size([2, 3, 4, 5])
print(x_rearranged.shape)  # 输出: torch.Size([2, 5, 3, 4])

2 reduce 计算平均值

(batch_size, height, width, channels) 的张量沿着通道维度求平均值，得到形状为 (batch_size, 1, height, width) 的输出。这可以使用 einops 的 reduce 函数来实现：

from einops import reduce# 创建一个形状为 (2, 3, 4, 5) 的张量
x = torch.randn(2, 3, 4, 5)# 使用 einops 沿着通道维度求平均值
x_mean = reduce(x, 'batch channels height width -> batch 1 height width', 'mean')print(x.shape)      # 输出: torch.Size([2, 3, 4, 5])
print(x_mean.shape) # 输出: torch.Size([2, 1, 4, 5])

在这个例子中，我们使用了 reduce 函数，并提供了描述输入和输出形状的字符串 'batch channels height width -> batch 1 height width' 和 'mean' 来实现沿着通道维度求平均值的操作。

3 repeat和stack

repeat 和 stack 的对比：

repeat 用于重复张量的元素，可以在指定的轴上重复张量的某些维度，但不会改变张量的维度数量。它通过重复元素来改变张量的形状。例如，将形状为 (batch_size, height, width, channels) 的张量在高度、宽度和通道维度上重复一次，形状变为 (batch_size, height * 2, width * 2, channels)。
stack 用于堆叠多个张量，可以在指定的轴上增加新的维度，将多个张量堆叠在一起形成一个新的维度。它通过堆叠张量来改变张量的维度数量。例如，将两个形状相同的张量在通道维度上堆叠，形状变为 (batch_size, 2, height, width, channels)。

因此，repeat 主要用于改变张量的形状，而 stack 主要用于改变张量的维度数量。两者在功能和使用场景上有明显的区别，需要根据具体的需求来选择使用哪个函数。

repeat 的示例：

import torch
from einops import repeat# 创建一个形状为 (2, 3) 的张量
x = torch.tensor([[1, 2, 3], [4, 5, 6]])# 使用 einops 在行和列上分别重复一次
x_repeated = repeat(x, 'h w -> h (repeat) w (repeat)', repeat=2)print(x.shape)           # 输出: torch.Size([2, 3])
print(x_repeated.shape)  # 输出: torch.Size([2, 2, 3, 2])

在这个例子中，我们使用 repeat 函数将形状为 (2, 3) 的张量 x 在行和列上分别重复一次，得到形状为 (2, 2, 3, 2) 的张量 x_repeated。

stack 的示例：

import torch
from einops import stack# 创建两个形状相同的张量
x = torch.tensor([[1, 2, 3], [4, 5, 6]])
y = torch.tensor([[7, 8, 9], [10, 11, 12]])# 使用 einops 在新的维度上堆叠张量
stacked_tensor = stack([x, y], 'b h w -> b (stack) h w')print(stacked_tensor.shape)  # 输出: torch.Size([2, 2, 3])