在Python中高效地进行张量收缩

Question

在Python中高效地进行张量收缩

3

我有一个张量列表L，每个张量(ndarray对象)都有多个指标。我需要根据连接的图表收缩这些指标。

连接通过元组的形式编码在列表中，格式为((m,i),(n,j))，表示“将张量L[m]的第i个指标与张量L[n]的第j个指标进行收缩”。

如何处理复杂的连接图？第一个问题是，一旦我收缩一对指标，结果就是一个新张量，它不属于列表L。但即使我解决了这个问题(例如，为所有张量的所有指标提供唯一标识符)，仍然存在一个问题，即可以任选顺序执行收缩，某些选择会在计算过程中产生不必要的巨大的临时结果（即使最终结果很小）。有什么建议吗？

- Ziofil

2个回答

1

可能有帮助：查看https://arxiv.org/abs/1402.0939，这是一个高效处理所谓张量网络问题的框架的描述，可以在单个函数ncon(...)中进行。据我所见，它的实现直接可用于Matlab（可以在链接中找到）和Python3（https://github.com/mhauru/ncon）。

- mmarah

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Andras Deak -- Слава Україні · Accepted Answer

除了内存考虑之外，我认为您可以在单个对einsum的调用中完成缩写，尽管您需要进行一些预处理。我不完全确定您所说的“当我收缩一对索引时，结果是一个新张量，它不属于列表L”是什么意思，但我认为在单个步骤中进行收缩将完全解决这个问题。

我建议使用einsum的替代数字索引语法：

einsum(op0, sublist0, op1, sublist1, ..., [sublistout])

所以你需要做的是将要缩小的索引编码为整数序列。首先，您需要初始化一系列唯一的索引，并保留另一个副本用作sublistout。然后，在遍历连接图时，您需要在必要时将缩小的索引设置为相同的索引，并同时从sublistout中删除缩小的索引。

import numpy as np

def contract_all(tensors,conns):
    '''
    Contract the tensors inside the list tensors
    according to the connectivities in conns

    Example input:
    tensors = [np.random.rand(2,3),np.random.rand(3,4,5),np.random.rand(3,4)]
    conns = [((0,1),(2,0)), ((1,1),(2,1))]
    returned shape in this case is (2,3,5)
    '''

    ndims = [t.ndim for t in tensors]
    totdims = sum(ndims)
    dims0 = np.arange(totdims)
    # keep track of sublistout throughout
    sublistout = set(dims0.tolist())
    # cut up the index array according to tensors
    # (throw away empty list at the end)
    inds = np.split(dims0,np.cumsum(ndims))[:-1]
    # we also need to convert to a list, otherwise einsum chokes
    inds = [ind.tolist() for ind in inds]

    # if there were no contractions, we'd call
    # np.einsum(*zip(tensors,inds),sublistout)

    # instead we need to loop over the connectivity graph
    # and manipulate the indices
    for (m,i),(n,j) in conns:
        # tensors[m][i] contracted with tensors[n][j]

        # remove the old indices from sublistout which is a set
        sublistout -= {inds[m][i],inds[n][j]}

        # contract the indices
        inds[n][j] = inds[m][i]

    # zip and flatten the tensors and indices
    args = [subarg for arg in zip(tensors,inds) for subarg in arg]

    # assuming there are no multiple contractions, we're done here
    return np.einsum(*args,sublistout)

一个简单的例子：

>>> tensors = [np.random.rand(2,3), np.random.rand(4,3)]
>>> conns = [((0,1),(1,1))]
>>> contract_all(tensors,conns)
array([[ 1.51970003,  1.06482209,  1.61478989,  1.86329518],
       [ 1.16334367,  0.60125945,  1.00275992,  1.43578448]])
>>> np.einsum('ij,kj',tensors[0],tensors[1])
array([[ 1.51970003,  1.06482209,  1.61478989,  1.86329518],
       [ 1.16334367,  0.60125945,  1.00275992,  1.43578448]])

如果有多个缩并，循环中的物流变得更加复杂，因为我们需要处理所有重复项。然而，逻辑是相同的。此外，上述内容显然缺少检查以确保相应的索引可以进行缩并。

回过头来我意识到默认的sublistout不必指定，einsum会自动使用该顺序。我决定在代码中保留该变量，因为如果我们想要一个非平凡的输出索引顺序，我们将不得不适当地处理该变量，并且它可能会派上用场。

关于缩并顺序的优化，从版本1.12开始（正如@hpaulj在现已删除的评论中所述），您可以在np.einsum中进行内部优化。此版本引入了optimize可选关键字参数到np.einsum中，允许选择一个缩并顺序以在计算时间上减少而在内存上增加。将'greedy'或'optimal'作为optimize关键字传递将使numpy选择大致按维度大小递减的缩并顺序。 optimize关键字的可用选项来自于显然未记录在在线文档中（但help()幸运地起作用）的函数np.einsum_path。

einsum_path(subscripts, *operands, optimize='greedy')

Evaluates the lowest cost contraction order for an einsum expression by
considering the creation of intermediate arrays.

< p>从np.einsum_path输出的收缩路径也可以用作np.einsum的optimize参数的输入。在您的问题中，您担心使用太多内存，因此我怀疑默认情况下没有优化（可能运行时间较长，内存占用较小）。