python中的高效张量收缩

时间:2017-02-03 23:11:35

标签: python numpy vectorization numpy-einsum

我有张量L张量(ndarray个对象),每个都有几个索引。我需要根据连接图收缩这些指数。

连接以((m,i),(n,j))形式的元组列表编码,表示“与 i - 张量L[m]的索引与 j < / em> - 张量L[n]的索引。

如何处理非平凡连接图?第一个问题是,只要我收缩一对索引,结果就是一个不属于列表L的新张量。但即使我解决了这个问题(例如,通过为所有张量的所有索引提供唯一标识符),也存在一个问题,即人们可以选择任何顺序来执行收缩,而某些选择会在计算中期产生不必要的巨大影响(即使最终结果很小)。建议?

2 个答案:

答案 0 :(得分:5)

除了内存注意事项,我相信你可以通过一次调用来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]])

如果有多次收缩,循环中的物流变得有点复杂,因为我们需要处理所有重复项。然而,逻辑是一样的。此外,上述显然缺少检查以确保相应的指数可以签约。

事后我意识到不必指定默认sublistouteinsum无论如何都会使用该订单。我决定在代码中保留该变量,因为如果我们想要一个非平凡的输出索引顺序,我们必须适当地处理该变量,它可能会派上用场。

对于收缩顺序的优化,您可以在版本1.12的np.einsum中实现内部优化(正如@hpaulj在现在删除的注释中所述)。此版本向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.

np.einsum_path的输出收缩路径也可以用作optimize np.einsum参数的输入。在您的问题中,您担心使用的内存过多,因此我怀疑没有优化的默认值(可能更长的运行时间和更小的内存占用量)。

答案 1 :(得分:1)

可能有用:请看一下https://arxiv.org/abs/1402.0939,这是一个有效的框架,用于解决在单个函数ncon(...)中收缩所谓的张量网络的问题。据我所知,它的实现可直接用于Matlab(可在链接中找到)和Python3(https://github.com/mhauru/ncon)。