我倾向于使用
einsum
因为这通常是最容易写的:
def fast(wxf, wyf, wzf):
return np.einsum('if,jf,kf->ijk', wxf, wyf, wzf)
def slow(wxf, wyf, wzf):
N = len(wxf)
W = np.zeros((N, N, N))
for i in range(N):
for j in range(N):
for k in range(N):
W[i,j,k] = np.sum([wxf[i,f]*wyf[j,f]*wzf[k,f] for f in range(N)])
return W
def gen_ws(N):
wxf = np.random.random((N,N))
wyf = np.random.random((N,N))
wzf = np.random.random((N,N))
return wxf, wyf, wzf
给予
>>> ws = gen_ws(25)
>>> via_slow = slow(*ws)
>>> via_fast = fast(*ws)
>>> np.allclose(via_slow, via_fast)
True
和
>>> ws = gen_ws(100)
>>> %timeit fast(*ws)
10 loops, best of 3: 91.6 ms per loop