4D矩阵乘法

4D matrix multiplication

我一直未能找到使用 numpy 编写以下 for 循环的好方法。现在的写法当然效率很低,对我的真实数据(形状 512*8*8*512)来说不可行,但我只是没能有效地使用内置的矩阵乘法函数。

import numpy as np
#Create pseudo weights, biases, input
weights = np.random.rand(10, 8, 8, 10)
biases = np.random.rand(10)
pseudo_input = np.random.rand(10, 8, 8)
output_loops = np.zeros((weights.shape[0],))

for n in range(10):
    output_loops[n] += biases[n]
    for x in range(8):
        for y in range(8):
            for f in range(10):
                output_loops[n] += weights[f, x, y, n] * pseudo_input[f,x,y]

简单地将相关的迭代器移植到np.einsum中的einsum字符串表示法-

np.einsum('fxyn,fxy->n', weights, pseudo_input) + biases

我们也可以使用np.tensordot-

np.tensordot(weights, pseudo_input, axes=((0,1,2),(0,1,2))) + biases

使用可信赖的 np.dot 进行一些额外的整形,使形状达到 2D1D -

pseudo_input.ravel().dot(weights.reshape(-1,weights.shape[-1])) + biases

另一个可能更慢的解决方案:

output_loops = (weights * pseudo_input[...,np.newaxis]).sum(axis=(0, 1, 2))