Сократить количество вычислений при поэлементном умножении многомерных массивов

В наличии имеется трехмерный массив. Пример представлен ниже. Нужно перемножить между собой элементы последнего измерения. Подробности и пояснения задачи в примере ниже в коде.

matrix = tensor([[[0.4804, 0.8351, 0.2127],
                  [0.5993, 0.7906, 0.7307]],

                  [[0.0767, 0.2068, 0.2575],
                  [0.0811, 0.5181, 0.8339]]]) # исходный массив

first_elements_of_last_dimension = matrix[:, :, 0]

>>> tensor([[0.4804, 0.5993],
            [0.0767, 0.0811]]) # матрица первых элементов последнего измерения

second_elements_of_last_dimension = matrix[:, :, 1]

>>> tensor([[0.8351, 0.7906],
            [0.2068, 0.5181]]) # матрица вторых элементов последнего измерения

third_elements_of_last_dimension = matrix[:, :, 2]

>>> tensor([[0.2127, 0.7307],
            [0.2575, 0.8339]]) # матрица третьих элементов последнего измерения
mult_matr = first_elements_of_last_dimension * second_element_of_last_dimension * third_element_of_last_dimension # поэлементное (декартово) произведение

mult_matr
>>> tensor([[0.0853, 0.3462],
            [0.0041, 0.0350]])

Мое решение заключалось в итерировании по диапазону range(количество элементов в последнем измерении) и перемножении единичного вектора на каждом шаге итерации (не самый быстрый способ):

product = 1
for i in range(matr.size()[2]):
    product *= matr[:,:,i]

Для того, чтобы сократить количество вычислений, нужно реализовать следующий алгоритм: Умножить первые элементы матрицы на второй элементы матрицы, одновременно умножить последний на предпоследний. Повторять до тех пор пока не схлопнется. То есть пусть у нас имеется 5 элементов последнего измерения [0,1,2,3,4]. Сначала res1 = matr[:,:, 0] * matr[:,:,1], там же res2 = matr[:,:,3] * matr[:,:,4]. Потом product = res1*res2*matr[:,:,2]. Можно ли как-то иначе сократить количество вычислений?


Ответы (2 шт):

Автор решения: Александр
import numpy as np
from timeit import timeit

# input array
arr = np.array([[[0.4804, 0.8351, 0.2127],
                  [0.5993, 0.7906, 0.7307]],

                  [[0.0767, 0.2068, 0.2575],
                  [0.0811, 0.5181, 0.8339]]])

# make product
def make_task(arr):
    return np.prod(np.split(arr, arr.shape[2], axis=2), axis=0)

print(timeit("make_task(arr)", number=10000, globals=globals()))
print(make_task(arr))

# output:
"""
0.9346408000000002
[[[0.08533142]
  [0.34621047]]

 [[0.00408435]
  [0.03503874]]]
"""

я уверен, что встроенные функции numpy будут работать быстрее вашего кода

я также нашёл некоторую информацию, как конвертировать numpy и torch: https://pytorch.org/tutorials/beginner/former_torchies/tensor_tutorial.html надеюсь, поможет

поправка 1:

я сделал так, чтобы в конечном результате вместо

[[[0.08533142]
  [0.34621047]]

 [[0.00408435]
  [0.03503874]]]

было

[[0.08533142 0.34621047]
 [0.00408435 0.03503874]]

вот код:

# make product
def make_task(arr):
    product = np.prod(np.dsplit(arr, arr.shape[2]), axis=0)
    return np.reshape(product, product[:2])
→ Ссылка
Автор решения: MaxU

Воспользуйтесь torch.prod(input, dim, keepdim=False, dtype=None):

In [5]: res = torch.prod(matrix, dim=2)

In [6]: res
Out[6]:
tensor([[0.0853, 0.3462],
        [0.0041, 0.0350]])

Если в результате нужен 3D тензор размерности (2, 2, 1):

In [11]: res = torch.prod(matrix, dim=2, keepdim=True)

In [12]: res
Out[12]:
tensor([[[0.0853],
         [0.3462]],

        [[0.0041],
         [0.0350]]])

сравнение скорости выполнения:

In [21]: t = torch.rand(1000, 1000, 100)

In [22]: t.shape
Out[22]: torch.Size([1000, 1000, 100])

In [23]: %timeit make_task(t.numpy())
1.77 s ± 68.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [24]: %timeit torch.prod(t, dim=2)
20.7 ms ± 669 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

разница в скорости:

In [25]: 1.77 * 1000 / 20.7
Out[25]: 85.5072463768116

UPD: сравнение скорости выполнения torch.prod() - GPU vs CPU:

In [1]: import torch

In [2]: t = torch.rand(1000, 1000, 100)

In [3]: a = t.to("cuda")

In [4]: %timeit torch.prod(t, dim=2)
21.1 ms ± 275 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [5]: %timeit torch.prod(a, dim=2)
2.44 ms ± 3.63 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
→ Ссылка