Faster definition of "matrix multiplication" in Python

user_11493659 0
Viewed 311

I need to define matrix multiplication from scratch, as instead of multiplying each constant together, each constant is actually another array and any two arrays need to be "convolved" together (I don't think it's necessary to define what a convolution is here).

I have made a picture that hopefully explains what I'm trying to say better:

here

The code I have to do this with is this:

for row in range(arr1.shape[2]):
    for column in range(arr2.shape[3]):
        for index in range(arr2.shape[2]): # Could also be "arr1.shape[3]"
            out[:, :, row, column] += convolve(
                arr2[:, :, :  , column][:, :, index],
                arr1[:, :, row, :     ][:, :, index]
            )

However, this method had proved to be very slow for me, so I was wondering if there was a faster way to do this.

matrix-multiplication convolution numpy python performance
1 Answers

If the intermediate fits in memory the following should be reasonably efficient

import numpy as np
from scipy.signal import fftconvolve,convolve

# example
rng = np.random.default_rng()
A = rng.random((5,6,2,3))                    
B = rng.random((4,3,3,4))                    

# custom matmul

Ae,Be = A[...,None],B[:,:,None]
shsh = np.maximum(Ae.shape[2:],Be.shape[2:])
Ae = np.broadcast_to(Ae,(*Ae.shape[:2],*shsh))
Be = np.broadcast_to(Be,(*Be.shape[:2],*shsh))
C = fftconvolve(Ae,Be,axes=(0,1),mode='valid').sum(3)         

# original loop for reference

out = np.zeros_like(C)
for row in range(A.shape[2]):
    for column in range(B.shape[3]):
        for index in range(B.shape[2]): # Could also be "A.shape[3]"
            out[:, :, row, column] += convolve(
                B[:, :, :  , column][:, :, index],
                A[:, :, row, :     ][:, :, index],
                mode='valid'
            )

print(np.allclose(C,out))

# True

By doing the convolution in bulk we reduce the total number of fft's we have to do.

If need be this could be further optimized for both speed and memory by doing the sum reduction in Fourier space using einsum. This would require doing the fft convolution by hand, though.

user_72073920
Related