Эффективное сокращение тензора в python

Question

У меня есть список L тензоров (ndarray объектов) с несколькими индексами каждый. Мне нужно сжать эти индексы согласно графику соединений.

подключений кодируются в списке кортежей в форме ((m,i),(n,j)) означающий «контракт I его индекса тензора L[m] с -й индексом тензора L[n]J.

Как может Я обрабатываю нетривиальные диаграммы связности? Первая проблема заключается в том, что, как только я сжимаю пару индексов, результатом является новый тензор, который не принадлежит списку L. Но даже если бы я решил это (например, предоставив уникальный идентификатор ко всем индексам всех тензоров), есть проблема, что можно выбрать любой порядок выполнения сокращений и некоторые варианты в средние вычисления приносят неоправданно огромные звери (даже если конечный результат мал). Предложения?

Answer 1

Вопросы памяти в сторону, я считаю, что вы можете сделать сокращения в один звонок до 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, позволяющий выбрать порядок сокращения, который сокращает время вычислений за счет стоимости памяти. Передача '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 также могут быть использованы в качестве вход для аргумента optimizenp.einsum. В вашем вопросе вас беспокоило слишком много памяти, поэтому я подозреваю, что по умолчанию нет оптимизации (с потенциально более продолжительным временем работы и меньшим объемом памяти).