Быстрое вращение тензора с NumPy
В основе приложения (написанного на Python и использующего NumPy), мне нужно повернуть тензор 4-го порядка. На самом деле, мне нужно много раз поворачивать много тензоров, и это мое узкое место. Моя наивная реализация (ниже) с восемью вложенными циклами кажется довольно медленной, но я не вижу возможности использовать операции с матрицей NumPy и, надеюсь, ускорить работу. У меня такое чувство, что я должен использовать np.tensordot
, но я не вижу, как.
Математически элементы повернутого тензора T 'задаются формулой: T' ijkl= & Sigma; g ia g jb g kc g ld T abcd с суммой повторяющиеся индексы справа. T и Tprime представляют собой 3 * 3 * 3 * 3 массивы NumPy, а матрица поворота g представляет собой массив 3 * 3 NumPy. Моя медленная реализация (принимающая ~ 0,04 секунды за звонок) ниже.
#!/usr/bin/env python
import numpy as np
def rotT(T, g):
Tprime = np.zeros((3,3,3,3))
for i in range(3):
for j in range(3):
for k in range(3):
for l in range(3):
for ii in range(3):
for jj in range(3):
for kk in range(3):
for ll in range(3):
gg = g[ii,i]*g[jj,j]*g[kk,k]*g[ll,l]
Tprime[i,j,k,l] = Tprime[i,j,k,l] + \
gg*T[ii,jj,kk,ll]
return Tprime
if __name__ == "__main__":
T = np.array([[[[ 4.66533067e+01, 5.84985000e-02, -5.37671310e-01],
[ 5.84985000e-02, 1.56722231e+01, 2.32831900e-02],
[ -5.37671310e-01, 2.32831900e-02, 1.33399259e+01]],
[[ 4.60051700e-02, 1.54658176e+01, 2.19568200e-02],
[ 1.54658176e+01, -5.18223500e-02, -1.52814920e-01],
[ 2.19568200e-02, -1.52814920e-01, -2.43874100e-02]],
[[ -5.35577630e-01, 1.95558600e-02, 1.31108757e+01],
[ 1.95558600e-02, -1.51342210e-01, -6.67615000e-03],
[ 1.31108757e+01, -6.67615000e-03, 6.90486240e-01]]],
[[[ 4.60051700e-02, 1.54658176e+01, 2.19568200e-02],
[ 1.54658176e+01, -5.18223500e-02, -1.52814920e-01],
[ 2.19568200e-02, -1.52814920e-01, -2.43874100e-02]],
[[ 1.57414726e+01, -3.86167500e-02, -1.55971950e-01],
[ -3.86167500e-02, 4.65601977e+01, -3.57741000e-02],
[ -1.55971950e-01, -3.57741000e-02, 1.34215636e+01]],
[[ 2.58256300e-02, -1.49072770e-01, -7.38843000e-03],
[ -1.49072770e-01, -3.63410500e-02, 1.32039847e+01],
[ -7.38843000e-03, 1.32039847e+01, 1.38172700e-02]]],
[[[ -5.35577630e-01, 1.95558600e-02, 1.31108757e+01],
[ 1.95558600e-02, -1.51342210e-01, -6.67615000e-03],
[ 1.31108757e+01, -6.67615000e-03, 6.90486240e-01]],
[[ 2.58256300e-02, -1.49072770e-01, -7.38843000e-03],
[ -1.49072770e-01, -3.63410500e-02, 1.32039847e+01],
[ -7.38843000e-03, 1.32039847e+01, 1.38172700e-02]],
[[ 1.33639532e+01, -1.26331100e-02, 6.84650400e-01],
[ -1.26331100e-02, 1.34222177e+01, 1.67851800e-02],
[ 6.84650400e-01, 1.67851800e-02, 4.89151396e+01]]]])
g = np.array([[ 0.79389393, 0.54184237, 0.27593346],
[-0.59925749, 0.62028664, 0.50609776],
[ 0.10306737, -0.56714313, 0.8171449 ]])
for i in range(100):
Tprime = rotT(T,g)
Есть ли способ ускорить это? Создание кода для других рангов тензора было бы полезным, но менее важно.
Ответы
Ответ 1
Чтобы использовать tensordot
, вычислите внешнее произведение тензоров g
:
def rotT(T, g):
gg = np.outer(g, g)
gggg = np.outer(gg, gg).reshape(4 * g.shape)
axes = ((0, 2, 4, 6), (0, 1, 2, 3))
return np.tensordot(gggg, T, axes)
В моей системе это примерно в семь раз быстрее, чем решение Sven. Если тензор g
не меняется часто, вы также можете кэшировать тензор gggg
. Если вы сделаете это и включите некоторые микрооптимизации (вставляя код tensordot
, никаких проверок, не генерических форм), вы все равно можете сделать это в два раза быстрее:
def rotT(T, gggg):
return np.dot(gggg.transpose((1, 3, 5, 7, 0, 2, 4, 6)).reshape((81, 81)),
T.reshape(81, 1)).reshape((3, 3, 3, 3))
Результаты timeit
на моем домашнем ноутбуке (500 итераций):
Your original code: 19.471129179
Sven code: 0.718412876129
My first code: 0.118047952652
My second code: 0.0690279006958
Числа на моем рабочем компьютере:
Your original code: 9.77922987938
Sven code: 0.137110948563
My first code: 0.0569641590118
My second code: 0.0308079719543
Ответ 2
Вот как это сделать с помощью одного цикла Python:
def rotT(T, g):
Tprime = T
for i in range(4):
slices = [None] * 4
slices[i] = slice(None)
slices *= 2
Tprime = g[slices].T * Tprime
return Tprime.sum(-1).sum(-1).sum(-1).sum(-1)
По общему признанию, это немного сложно понять на первый взгляд, но это довольно быстро:)
Ответ 3
Благодаря тяжелой работе М. Вибе, следующая версия Numpy (которая, вероятно, будет 1.6), сделает это еще проще:
>>> Trot = np.einsum('ai,bj,ck,dl,abcd->ijkl', g, g, g, g, T)
Филипп подход в настоящее время в 3 раза быстрее, хотя, но, возможно, есть некоторые возможности для улучшения. Разница скоростей, вероятно, в основном связана с тем, что tensdord способен развернуть всю операцию как единый матричный продукт, который может быть передан BLAS, и, таким образом, избежать значительной части служебных данных, связанных с небольшими массивами, - это невозможно для общего Эйнштейна суммирование, так как не все операции, которые могут быть выражены в этой форме, решаются на один матричный продукт.
Ответ 4
Из любопытства я сравнил Cython реализацию наивного кода из вопрос с кодом numpy из @Philipp answer. Код Cython в четыре раза быстрее на моей машине:
#cython: boundscheck=False, wraparound=False
import numpy as np
cimport numpy as np
def rotT(np.ndarray[np.float64_t, ndim=4] T,
np.ndarray[np.float64_t, ndim=2] g):
cdef np.ndarray[np.float64_t, ndim=4] Tprime
cdef Py_ssize_t i, j, k, l, ii, jj, kk, ll
cdef np.float64_t gg
Tprime = np.zeros((3,3,3,3), dtype=T.dtype)
for i in range(3):
for j in range(3):
for k in range(3):
for l in range(3):
for ii in range(3):
for jj in range(3):
for kk in range(3):
for ll in range(3):
gg = g[ii,i]*g[jj,j]*g[kk,k]*g[ll,l]
Tprime[i,j,k,l] = Tprime[i,j,k,l] + \
gg*T[ii,jj,kk,ll]
return Tprime
Ответ 5
Я думал, что внесет относительно новые данные в эти тесты, используя parakeet, один из numpy
-aware JIT, которые возникли в последние несколько месяцев. (Другой, о котором я знаю, numba, но я не тестировал его здесь.)
После того, как вы выполните несколько лабораторных процедур установки для LLVM, вы можете украсить многие функции pure- numpy
(часто) ускорить их производительность:
import numpy as np
import parakeet
@parakeet.jit
def rotT(T, g):
# ...
Я только попытался применить JIT к коду Andrew в исходном вопросе, но он довольно хорошо ( > 10x speedup) для того, чтобы не писать какой-либо новый код вообще:
andrew 10 loops, best of 3: 206 msec per loop
andrew_jit 10 loops, best of 3: 13.3 msec per loop
sven 100 loops, best of 3: 2.39 msec per loop
philipp 1000 loops, best of 3: 0.879 msec per loop
Для этих таймингов (на моем ноутбуке) я запускал каждую функцию десять раз, чтобы дать JIT возможность идентифицировать и оптимизировать пути для горячих кодов.
Интересно, что предложения Свена и Филиппа все еще на порядок выше!
Ответ 6
Предполагаемый подход и код решения
Для эффективности памяти и, следовательно, эффективности работы мы можем использовать тензорное умножение матрицы по этапам.
Чтобы проиллюстрировать предпринятые шаги, позвольте использовать простейшее из решений с np.einsum
by @pv. -
np.einsum('ai,bj,ck,dl,abcd->ijkl', g, g, g, g, T)
Как видно, мы теряем первое измерение от g
с тензорным умножением между его четырьмя вариантами и T
.
Давайте сделаем эти суммы-сокращения для умножения тензорных матриц по шагам. Пусть начнется первый вариант g
и T
:
p1 = np.einsum('abcd, ai->bcdi', T, g)
Таким образом, мы получаем тензор размеров как обозначение строки: bcdi
. Следующие шаги включают суммирование, уменьшающее этот тензор до остальных трех вариантов g
, используемых в оригинальной имплантации einsum
. Следовательно, следующим сокращением будет -
p2 = np.einsum('bcdi, bj->cdij', p1, g)
Как видно, мы потеряли первые два измерения со строковыми обозначениями: a
, b
. Мы продолжаем его еще два шага, чтобы избавиться от c
и d
тоже и оставим с ijkl
в качестве окончательного вывода, например:
p3 = np.einsum('cdij, ck->dijk', p2, g)
p4 = np.einsum('dijk, dl->ijkl', p3, g)
Теперь мы можем использовать np.tensordot
для этих суммирующих сокращений, что было бы намного более эффективным.
Окончательная реализация
Таким образом, перенося на np.tensordot
, мы бы выполнили окончательную реализацию так:
p1 = np.tensordot(T,g,axes=((0),(0)))
p2 = np.tensordot(p1,g,axes=((0),(0)))
p3 = np.tensordot(p2,g,axes=((0),(0)))
out = np.tensordot(p3,g,axes=((0),(0)))
Тест времени выполнения
Проведите тестирование всех подходов, основанных на NumPy, размещенных по другим сообщениям, чтобы решить проблему с производительностью.
Подходы как функции -
def rotT_Philipp(T, g): # @Philipp soln
gg = np.outer(g, g)
gggg = np.outer(gg, gg).reshape(4 * g.shape)
axes = ((0, 2, 4, 6), (0, 1, 2, 3))
return np.tensordot(gggg, T, axes)
def rotT_Sven(T, g): # @Sven Marnach soln
Tprime = T
for i in range(4):
slices = [None] * 4
slices[i] = slice(None)
slices *= 2
Tprime = g[slices].T * Tprime
return Tprime.sum(-1).sum(-1).sum(-1).sum(-1)
def rotT_pv(T, g): # @pv. soln
return np.einsum('ai,bj,ck,dl,abcd->ijkl', g, g, g, g, T)
def rotT_Divakar(T,g): # Posted in this post
p1 = np.tensordot(T,g,axes=((0),(0)))
p2 = np.tensordot(p1,g,axes=((0),(0)))
p3 = np.tensordot(p2,g,axes=((0),(0)))
p4 = np.tensordot(p3,g,axes=((0),(0)))
return p4
Сроки с исходными размерами данных -
In [304]: # Setup inputs
...: T = np.random.rand(3,3,3,3)
...: g = np.random.rand(3,3)
...:
In [305]: %timeit rotT(T, g)
...: %timeit rotT_pv(T, g)
...: %timeit rotT_Sven(T, g)
...: %timeit rotT_Philipp(T, g)
...: %timeit rotT_Divakar(T, g)
...:
100 loops, best of 3: 6.51 ms per loop
1000 loops, best of 3: 247 µs per loop
10000 loops, best of 3: 137 µs per loop
10000 loops, best of 3: 41.6 µs per loop
10000 loops, best of 3: 28.3 µs per loop
In [306]: 6510.0/28.3 # Speedup with the proposed soln over original code
Out[306]: 230.03533568904592
Как обсуждалось в начале этого сообщения, мы пытаемся добиться эффективности памяти и, следовательно, повышения производительности. Позвольте проверить это, поскольку мы увеличиваем размеры набора данных -
In [307]: # Setup inputs
...: T = np.random.rand(5,5,5,5)
...: g = np.random.rand(5,5)
...:
In [308]: %timeit rotT(T, g)
...: %timeit rotT_pv(T, g)
...: %timeit rotT_Sven(T, g)
...: %timeit rotT_Philipp(T, g)
...: %timeit rotT_Divakar(T, g)
...:
100 loops, best of 3: 6.54 ms per loop
100 loops, best of 3: 7.17 ms per loop
100 loops, best of 3: 2.7 ms per loop
1000 loops, best of 3: 1.47 ms per loop
10000 loops, best of 3: 39.9 µs per loop
Ответ 7
Не новый ответ, так как все предыдущие хорошо справляются с вопросом. Больше похоже на комментарий, но я публикую его как ответ, чтобы иметь место для кода.
Хотя все ответы воспроизводят результат оригинального поста, я почти уверен, что код, указанный в оригинальном посте, неверен. Глядя на формулу T 'ijkl= & Sigma; g ia g jb g kc g ld T abcd, что, на мой взгляд, правильно, показатели для g, которые различаются при расчете каждой записи T ', представляют собой a, b, c & д. Однако в исходном предоставленном коде индексы, используемые для доступа к значениям g при вычислении gg, меняются местами в соответствии с формулой. Следовательно, я считаю, что следующий код действительно обеспечивает правильную реализацию формулы:
def rotT(T, g):
Tprime = np.zeros((3, 3, 3, 3))
for i in range(3):
for j in range(3):
for k in range(3):
for l in range(3):
for a in range(3):
for b in range(3):
for c in range(3):
for d in range(3):
Tprime[i, j, k, l] += \
g[i, a] * g[j, b] * \
g[k, c] * g[l, d] * T[a, b, c, d]
Эквивалентные, но более быстрые вызовы обновлений einsum и tenordot на:
Tprime = np.tensordot(g, np.tensordot(g, np.tensordot(
g, np.tensordot(g, T, (1, 3)), (1, 3)), (1, 3)), (1, 3))
Tprime = np.einsum('ia, jb, kc, ld, abcd->ijkl', g, g, g, g, T)
Кроме того, использование @jit(nopython=True)
из numba для функции наивных циклов в пять раз быстрее, чем использование numpy.tensordot
на моей машине.