python-einops

what is einops

einops 是一个用于张量操作的 Python 库，它的核心思想是使用一种更直观、更声明性的方式来重塑、转置和执行其他多维数组操作。它的名字是 “Einstein operations” 的缩写，因为它受到爱因斯坦求和约定（Einstein summation notation）的启发。

einops 的用法

传统上，在 NumPy、PyTorch、TensorFlow 等库中进行张量操作，你需要使用一系列函数，比如 reshape, transpose, permute, view 等。这很容易出错，特别是当张量的维度很多时，你可能会混淆轴的顺序和大小。

einops 通过一个统一的、高可读性的表达式来解决这个问题。它的核心函数是 rearrange, reduce 和 repeat。

  1. rearrange
    这是 einops 最常用的函数，它结合了 重塑 (reshape) 和 转置 (transpose) 的功能。它的语法非常直观：einops.rearrange(tensor, ‘… -> …’)。

    1
    2
    3
    4
    import einops
    import torch
    tensor = torch.randn(2, 3, 4)
    print(einops.rearrange(tensor, 'b c h w -> b h w c'))

    另一个例子是，将形状为 (batch, channels, height, width) 的图像分割成不重叠的小块（patch），并展平。

    1
    2
    # 假设 patch 大小是 16x16
    print(einops.rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=16, p2=16))

  2. reduce
    这个函数用于执行各种常见的聚合操作，比如求和、求平均、最大值、最小值等。它的语法是：einops.reduce(tensor, ‘… -> …’, ‘op’)。

    1
    2
    from einops import reduce
    output_tensor = reduce(input_tensor, 'b c h w -> c', 'mean')

    这个表达式表示将 b, h, w 维度求平均，只保留 c 维度。

  3. repeat
    这个函数用于重复张量的某些维度，它的语法是：einops.repeat(tensor, ‘… -> …’, **kwargs)。

    1
    2
    from einops import repeat
    output_tensor = repeat(img, 'h w -> h w 3')

    将一个形状为 (height, width) 的灰度图像复制成 3 个通道。