当前位置：首页 > article >正文

深度学习张量变换操作利器 einops 基础实践

article 2024/9/20 17:39:55

今天在一个项目调试的时候无意间报错：

以前其实并没有怎么多接触过einops，今天正好碰到了，就简单总结记录下。

解决上面的报错很简单，直接pip安装即可：

einops 是一个用于操作张量的库，它提供了一种简洁且直观的方式来处理复杂的张量操作。einops 的设计目标是使张量操作更加易读、易写和易维护。它的名字来源于 "Einstein-Inspired Notation for Operations on Tensors"，即受爱因斯坦启发用于张量操作的符号。

主要特点

简洁的语法：einops 使用了一种类似于数学表达式的语法，使得张量操作的代码更加直观和易读。
灵活的操作：einops 支持多种张量操作，包括重塑（reshape）、转置（transpose）、拆分（split）、合并（merge）等。
跨框架支持：einops 可以在多个深度学习框架中使用，包括 PyTorch、TensorFlow、JAX 等。
性能优化：einops 的实现经过了优化，能够在不牺牲性能的情况下提供简洁的语法。

核心功能

1. `rearrange`

rearrange 用于重塑张量的形状。它允许你通过指定新的维度顺序和形状来重新排列张量的元素。

import torch
from einops import rearrange

# 创建一个形状为 (2, 3, 4) 的张量
x = torch.arange(24).reshape(2, 3, 4)

# 将张量重新排列为 (3, 2, 4)
y = rearrange(x, 'a b c -> b a c')
print(y.shape)  # 输出: torch.Size([3, 2, 4])

2. `reduce`

reduce 用于对张量进行降维操作，例如求和、平均、最大值、最小值等。

from einops import reduce

# 创建一个形状为 (2, 3, 4) 的张量
x = torch.arange(24).reshape(2, 3, 4)

# 对最后一个维度求和
y = reduce(x, 'a b c -> a b', 'sum')
print(y.shape)  # 输出: torch.Size([2, 3])

3. `repeat`

repeat 用于重复张量的某些维度。

from einops import repeat

# 创建一个形状为 (2, 3) 的张量
x = torch.arange(6).reshape(2, 3)

# 在第一个维度上重复 2 次
y = repeat(x, 'a b -> (2 a) b')
print(y.shape)  # 输出: torch.Size([4, 3])

4. `pack` 和 `unpack`

pack 和 unpack 用于将多个张量打包成一个张量，或者将一个张量解包成多个张量。

安装

你可以通过 pip 安装 einops：

pip install einops

使用示例

以下是一个使用 einops 进行复杂张量操作的示例：

import torch
from einops import rearrange, reduce, repeat

# 创建一个形状为 (2, 3, 4) 的张量
x = torch.arange(24).reshape(2, 3, 4)

# 重新排列维度
y = rearrange(x, 'a b c -> b a c')
print(y.shape)  # 输出: torch.Size([3, 2, 4])

# 对最后一个维度求和
z = reduce(y, 'a b c -> a b', 'sum')
print(z.shape)  # 输出: torch.Size([3, 2])

# 在第一个维度上重复 2 次
w = repeat(z, 'a b -> (2 a) b')
print(w.shape)  # 输出: torch.Size([6, 2])