Einops（整理版）
原始资料: Learn Einops - ReCoDE Deep Learning Best Practices created: 2026-07-02 22:12 整理说明: 本版本结合原始笔记和可读取的原始教程重排、补全和翻译，保留常用英文术语。
内容简要概括
Einops 用可读的 pattern 表达 tensor 的 reshape、transpose、reduce 和 repeat，让维度变换比连续调用 reshape、transpose 更直观。最需要记住的三个函数是 rearrange、reduce、repeat：分别用于改变维度结构、压缩维度、复制或扩展维度。它可以和 NumPy、PyTorch、TensorFlow 配合使用，也提供可放进 PyTorch 模型里的 Rearrange layer。
Einops、rearrange、reduce、repeat、Rearrange、tensor、reshape、transpose、batch、pattern、axis、NumPy、PyTorch
目录
1. 先记住三个函数
原笔记的核心记法很好，保留并补充：
rearrange：改变维度结构，比如 reshape、transpose、flatten、split、combine
reduce：压缩某些维度，比如 mean、sum、max
repeat：复制或扩展维度，比如增加 channel 或 tile
三者共同特点是用字符串 pattern 描述输入维度和输出维度：
"输入维度表达式 -> 输出维度表达式"
例如：
rearrange(x, "h w -> w h")
表示把二维数组从 height-width 顺序换成 width-height 顺序。
2. 安装与基本语法
安装：
pip install einops
常用导入：
from einops import rearrange, reduce, repeat
pattern 中常见字段：
| 字段 | 说明 |
|---|---|
b |
batch size。 |
c |
channel。 |
h |
height。 |
w |
width。 |
(h w) |
把多个维度合并成一个维度，或把一个维度拆成多个维度。 |
h=2 |
显式指定某个维度大小，常用于拆分维度时消除歧义。 |
3. rearrange：重排和合并维度
准备一个 4x4 数组：
import numpy as np
from einops import rearrange
x = np.array([
[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12],
[13, 14, 15, 16],
]).astype(np.float32)
把 4x4 变成 2x8：
x_reshaped = rearrange(x, "(h w1) w2 -> h (w1 w2)", h=2, w1=2, w2=4)
print(x_reshaped)
参数说明：
| 参数 | 作用 |
|---|---|
"(h w1) w2 -> h (w1 w2)" |
左侧表示输入维度，右侧表示输出维度。 |
(h w1) |
把输入第一个维度拆成 h 和 w1。 |
h (w1 w2) |
输出中保留 h，并把 w1 和 w2 合并。 |
h=2, w1=2, w2=4 |
指定拆分和合并时需要的具体维度大小。 |
转置二维数组：
x_transposed = rearrange(x, "h w -> w h")
print(x_transposed)
把 batch 中每张二维图 flatten：
batch = np.array([x, x])
batch_processed = rearrange(batch, "b h w -> b (h w)")
print(batch_processed.shape)
这里 b h w -> b (h w) 表示保留 batch 维，把每个样本内部的 h 和 w 合并。
4. reduce：压缩维度
对每一行求平均：
from einops import reduce
x_mean = reduce(x, "h w -> h", "mean")
print(x_mean)
参数说明：
| 参数 | 作用 |
|---|---|
"h w -> h" |
输入是二维，输出只保留 h，因此 w 维会被压缩。 |
"mean" |
压缩方式是求平均。也可以按任务使用 "sum"、"max" 等。 |
记法：在输出侧消失的维度，就是被 reduce 的维度。
5. repeat：复制或扩展维度
给二维数组增加一个 channel 维，并复制 3 份：
from einops import repeat
x_repeated = repeat(x, "h w -> h w c", c=3)
print(x_repeated.shape)
参数说明：
| 参数 | 作用 |
|---|---|
"h w -> h w c" |
输入没有 c，输出新增 c 维。 |
c=3 |
新增维度大小为 3。 |
repeat 适合构造重复通道、扩展 batch、tile 特征等操作。实际训练中要注意内存开销，因为复制会增加 tensor 大小。
6. 批量数据与 PyTorch layer
Einops 可以直接处理 batch 维度：
batch_processed = rearrange(batch, "b h w -> b (h w)")
在 PyTorch 模型中，可以使用 layer 形式：
import torch
from einops.layers.torch import Rearrange
x_torch = torch.tensor(x)
layer = Rearrange("h w -> h w 1")
x_torch_reshaped = layer(x_torch)
print(x_torch_reshaped.shape)
Rearrange 适合放进 nn.Sequential，让模型结构中的维度变换也保持可读：
from torch import nn
from einops.layers.torch import Rearrange
model = nn.Sequential(
Rearrange("b c h w -> b (c h w)"),
nn.Linear(28 * 28, 10),
)
阅读 Einops pattern 时，先看箭头左边的输入维度，再看箭头右边哪些维度被保留、交换、合并、拆分、新增或压缩。