NumPy einsum 详解：理解张量运算的精髓-Python教程-PHP中文网

NumPy einsum 详解：理解张量运算的精髓

花韻仙語

发布： 2025-10-23 09:14:01

原创

149人浏览过

numpy einsum 详解：理解张量运算的精髓

本文旨在深入解析 NumPy 中 `einsum` 函数的用法，通过具体示例和代码演示，帮助读者理解其在张量运算中的作用，并掌握利用 `einsum` 实现高效、灵活的张量操作的方法。文章将从基本概念入手，逐步剖析 `einsum` 的运算规则，并提供等效的循环实现，以便读者更好地理解其内部机制。

np.einsum 是 NumPy 中一个强大的函数，用于执行爱因斯坦求和约定。它允许用户以简洁而灵活的方式表达各种张量运算，例如矩阵乘法、张量缩并、求迹等。理解 einsum 的工作原理对于高效处理多维数组至关重要。

einsum 的基本语法

einsum 的基本语法如下：

np.einsum(subscripts, *operands, out=None, dtype=None, order='K', casting='safe', optimize=False)

登录后复制

其中，subscripts 是一个字符串，用于指定张量的维度以及如何进行求和。operands 则是要进行运算的张量。

理解 subscripts 字符串

subscripts 字符串的核心在于定义了输入张量的维度标签，以及输出张量的维度标签。相同的维度标签表示在该维度上进行求和（缩并）。

例如，'ijk,jil->kl' 表示：

第一个输入张量 a 的维度是 ijk。
第二个输入张量 b 的维度是 jil。
输出张量的维度是 kl。
由于 i 和 j 在输入张量中出现，但在输出张量中没有出现，因此在 i 和 j 维度上进行求和。

示例分析

考虑以下示例：

import numpy as np

a = np.arange(8.).reshape(4, 2, 1)
b = np.arange(16.).reshape(2, 4, 2)

result = np.einsum('ijk,jil->kl', a, b)
print(result)

登录后复制

为了更好地理解 einsum 的运算过程，我们可以将其分解为更细粒度的步骤。首先，考虑不进行任何求和的情况：

商汤商量

商汤科技研发的AI对话工具，商量商量，都能解决。

查看详情

result_no_sum = np.einsum('ijk,jil->ijkl', a, b)
print(result_no_sum)

登录后复制

这个操作将生成一个形状为 (4, 2, 1, 2) 的张量，其中每个元素是 a 和 b 中对应元素的乘积，而没有进行任何求和。

接下来，为了得到原始 einsum 的结果，我们需要在轴 0 和轴 1 上进行求和：

result_sum_1 = result_no_sum.sum(axis=1)
print(result_sum_1)

result_sum_0 = result_sum_1.sum(axis=0)
print(result_sum_0)

登录后复制

result_sum_0 的结果与 np.einsum('ijk,jil->kl', a, b) 的结果相同。

等效的循环实现

为了更深入地理解 einsum 的运算过程，我们可以使用循环来实现相同的功能：

def sum_array(A, B):
    i_len, j_len, k_len = A.shape
    _, _, l_len = B.shape

    ret = np.zeros((k_len, l_len))
    for i in range(i_len):
        for j in range(j_len):
            for k in range(k_len):
                for l in range(l_len):
                    ret[k, l] += A[i, j, k] * B[j, i, l]
    return ret

result_loop = sum_array(a, b)
print(result_loop)

登录后复制

这段代码的逻辑与 einsum('ijk,jil->kl', a, b) 完全相同。通过显式循环，我们可以清楚地看到每个元素的乘积以及求和的过程。