优化Einsum

优化Einsum：张量 contraction 顺序优化器

优化Einsum可以显著减少类似einsum的表达式（例如，np.einsum、dask.array.einsum、pytorch.einsum、tensorflow.einsum）的整体执行时间，通过优化表达式的 contraction 顺序并将许多操作调度到标准BLAS、cuBLAS或其他专用程序。

优化Einsum对后端不可知，可以处理NumPy、Dask、PyTorch、Tensorflow、CuPy、Sparse、Theano、JAX和Autograd数组，以及可能符合标准API的任何库。更多信息请参阅文档。

示例用法

函数 opt_einsum.contract 通常可以作为 einsum 函数的直接替代品，无需对代码进行修改即可提供更优的性能。在此示例中，我们将展示带优化和不带优化的张量合同运算。

import numpy as np
from opt_einsum import contract

N = 10
C = np.random.rand(N, N)
I = np.random.rand(N, N, N, N)

%timeit np.einsum('pi,qj,ijkl,rk,sl->pqrs', C, C, I, C, C)
1 loops, best of 3: 934 ms per loop

%timeit contract('pi,qj,ijkl,rk,sl->pqrs', C, C, I, C, C)
1000 loops, best of 3: 324 us per loop

在这个特定示例中，我们看到了大约 3000 倍的性能提升，与未优化的合同运算相比，这种情况并不罕见。有关使用其他后端的更多信息，请参阅后端示例。

特性

此存储库中发现的算法通常为上述许多项目的 einsum 优化提供动力。例如，np.einsum 的优化已传入上游，并且此存储库中可以找到的大多数功能都可以通过 np.einsum(..., optimize=True) 启用。然而，此存储库通常具有更复杂的合同优化算法。

opt_einsum 启用了以下功能

• 检查选择的路径的详细信息。
• 使用多个后端执行合同，包括 GPU 和 TensorFlow、PyTorch 等库。
• 生成可重用的表达式，可能包含常量张量，这些表达式可以编译以获得更好的性能。
• 使用任意数量的索引为数百甚至数千个张量找到合同。
• 在多个合同之间共享中间计算。
• 使用 autograd 或 jax 计算张量合同的梯度。

有关更多功能，请参阅文档！

安装

可以通过 pip install opt_einsum 或从 conda conda install opt_einsum -c conda-forge 安装 opt_einsum。有关其他安装方法的更多信息，请参阅安装文档。

引用

如果此代码对您的科研有益，请通过引用支持我们。

Daniel G. A. Smith 和 Johnnie Gray，opt_einsum - 用于优化 einsum-like 表达式的合同顺序的 Python 包。《开源软件杂志》，2018 年，3(26)，753

DOI: https://doi.org/10.21105/joss.00753

贡献

欢迎所有贡献，包括错误报告、错误修复、文档改进、增强和想法。

有关如何贡献的详细概述，请参阅贡献指南。