除了内存考虑之外,我认为您可以在单个对
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)
sublistout = set(dims0.tolist())
inds = np.split(dims0,np.cumsum(ndims))[:-1]
inds = [ind.tolist() for ind in inds]
for (m,i),(n,j) in conns:
sublistout -= {inds[m][i],inds[n][j]}
inds[n][j] = inds[m][i]
args = [subarg for arg in zip(tensors,inds) for subarg in arg]
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
参数的输入。在您的问题中,您担心使用太多内存,因此我怀疑默认情况下没有优化(可能运行时间较长,内存占用较小)。
optimize='optimal'
可以让我计算更大的einsum
数组,链接:https://dev59.com/Rp7ha4cB1Zd3GeqPm66-。迄今为止这是我使用新功能的唯一经验。 - hpauljoptimization=True
我得到了约1000倍的性能提升! - Ziofiloptimize=True
,我无法再编辑评论了。顺便说一下,非常感谢你! - Ziofil