Расширенное вещание в TensorFlow (или Numpy) - PullRequest
0 голосов
/ 24 июня 2019

В TensorFlow у меня есть тензор ранга 2 M (матрица) формы [D, D] и тензор ранга 3 T фигуры [D, D, D].

Мне нужно объединить ихчтобы сформировать новую матрицу R следующим образом: элемент R[a, b+c-a] задается суммой всех элементов T[a, b, c]*M[b, c], где b+c-a является константой (где b+c-a должно быть между 0 и D-1).

Неэффективный способ создания R - это использование вложенных циклов for над индексами и проверка того, что b+c-a не превышает D-1 (например, в numpy):

R = np.zeros([D,D])

for a in range(D):
  for b in range(D):
    for c in range(D):
      if 0 <= b+c-a < D:
        R[a, b+c-a] += T[a, b, c]*M[b, c]

но я бы хотел использовать широковещательную и / или другие более эффективные методы.

Как мне этого добиться?

1 Ответ

2 голосов
/ 24 июня 2019

Вы можете векторизовать этот расчет следующим образом:

import numpy as np

np.random.seed(0)
D = 10
M = np.random.rand(D, D)
T = np.random.rand(D, D, D)
# Original calculation
R = np.zeros([D, D])
for a in range(D):
    for b in range(D):
        for c in range(D):
            if 0 <= b + c - a < D:
                R[a, b + c - a] += T[a, b, c] * M[b, c]
# Vectorized calculation
tm = T * M
a = np.arange(D)[:, np.newaxis, np.newaxis]
b, c = np.ogrid[:D, :D]
col_idx = b + c - a
m = (col_idx >= 0) & (col_idx < D)
row_idx = np.tile(a, [1, D, D])
R2 = np.zeros([D, D])
np.add.at(R2, (row_idx[m], col_idx[m]), tm[m])
# Check result
print(np.allclose(R, R2))
# True

В качестве альтернативы, вы можете рассмотреть возможность использования Numba для ускорения циклов:

import numpy as np
import numba as nb

@nb.njit
def calculation_nb(T, M, D):
    tm = T * M
    R = np.zeros((D, D), dtype=tm.dtype)
    for a in nb.prange(D):
      for b in range(D):
        for c in range(max(a - b, 0), min(D + a - b, D)):
          R[a, b + c - a] += tm[a, b, c]
    return R

print(np.allclose(R, calculation_nb(T, M, D)))
# True

В нескольких быстрых тестах, даже безраспараллеливание, это намного быстрее, чем NumPy.

Добро пожаловать на сайт PullRequest, где вы можете задавать вопросы и получать ответы от других членов сообщества.
...