Векторизация Numpy 3D и 2D массива - PullRequest
2 голосов
/ 21 апреля 2019

Я пытаюсь создать K MxN matrices в Python, сохраненный в массиве (M,N,K) numpy, C, из двух матриц, A и B, с формами (K, M) и (K,N) соответственно.Первая матрица вычисляется как C0 = a0.T x b0, где a0 - первая строка A, а b1 - первая строка B, вторая матрица - C1 = a1.T x b0 и т. Д.

Прямо сейчас я использую цикл for для вычисления матриц.

import numpy as np
A = np.random.random((10,800))
B = np.random.random((10,500))
C = np.zeros((800,500,10))
for k in range(10):
    C[:,:,k] = A[k,:][:,None] @ B[k,:][None,:]

Поскольку операции независимы, мне было интересно, есть ли какой-нибудь питонный способ избежать цикла for.Возможно, я смогу векторизовать код, но не понимаю, как это можно сделать.

1 Ответ

0 голосов
/ 21 апреля 2019
In [235]: A = np.random.random((10,800)) 
     ...: B = np.random.random((10,500)) 
     ...: C = np.zeros((800,500,10)) 
     ...: for k in range(10): 
     ...:     C[:,:,k] = A[k,:][:,None] @ B[k,:][None,:] 
     ...:                                                                            
In [236]: C.shape                                                                    
Out[236]: (800, 500, 10)

Пакетное матричное произведение с последующей транспонированием

In [237]: np.allclose((A[:,:,None]@B[:,None,:]).transpose(1,2,0), C)                 
Out[237]: True

Но так как ось матричного произведения имеет размер 1, и нет другого суммирования, переданное умножение так же хорошо:

In [238]: np.allclose((A[:,:,None]*B[:,None,:]).transpose(1,2,0), C)                 
Out[238]: True

Время выполнения примерно одинаково

...