Skip to main content
 首页 » 编程设计

python之python中的高效张量收缩

2025年05月04日28langtianya

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

连接以 ((m,i),(n,j)) 形式编码在元组列表中,表示“收缩第 i 个索引张量 L[m] 与张量 L[n] 的第 j 索引。

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

请您参考如下方法:

抛开内存方面的考虑,我相信您可以在对 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 无论如何都会使用该顺序。我决定将该变量保留在代码中,因为万一我们需要一个重要的输出索引顺序,我们必须适本地处理该变量,它可能会派上用场。


至于收缩顺序的优化,您可以在 np.einsum 1.12 版中实现内部优化(正如@hpaulj 在现已删除的评论中指出的那样)。此版本向 np.einsum 引入了 optimize 可选关键字参数,允许选择以内存为代价减少计算时间的收缩顺序。传递 '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 的输出收缩路径也可以用作np.einsum 的optimize 参数的输入。在您的问题中,您担心使用的内存过多,因此我怀疑默认情况下没有优化(运行时间可能更长,内存占用量更小)。