Самый питонский способ умножения этих двух векторов?

Question

У меня есть два вектора с размерами:

A = (32,512,640)

B = (4,512)

Мне нужно умножить A и B так, чтобы я получил новый вектор:

C = (4,32,512,640)

Еще один способ думать об этомчто каждая строка вектора B умножается вдоль оси = -2 A, в результате чего получается новый 1,32,512,640 куб.Каждый ряд B можно зациклить, образуя 1,32,512,640 кубов, которые затем можно использовать для построения C, используя np.concatenate или np.vstack, например:

# Sample vectors, where the dimensions aren't necessarily known
a = np.arange(32*512*465, dtype='f4').reshape((32,512,465))
b = np.ones((4,512), dtype='f4')

# Using a loop
d = []
for row in b:
    d.append(np.expand_dims(row[None,:,None]*a, axis=0))

# Or using list comprehension
d = [np.expand_dims(row[None,:,None]*a,axis=0) for row in b]

# Stacking the final list
result = np.vstack(d)

Но мне интересно,можно использовать что-то вроде np.einsum или np.tensordot, чтобы векторизовать все в одну строку.Я все еще учусь использовать эти два метода, поэтому я не уверен, уместно ли это здесь.

Спасибо!

Answer 1

Мы можем использовать broadcasting после расширения размеров B с None/np.newaxis -

C = A * B[:,None,:,None]

С einsum, это было бы -

C = np.einsum('ijk,lj->lijk',A,B)

Здесь не происходит сокращения суммы, поэтому einsum будет не лучше, чем explicit-broadcasting.Но поскольку мы ищем решение Pythonic, которое можно было бы использовать, как только мы пройдем через строковую нотацию.

Давайте дадим немного времени для завершения -

In [15]: m,n,r,p = 32,512,640,4
    ...: A = np.random.rand(m,n,r)
    ...: B = np.random.rand(p,n)

In [16]: %timeit A * B[:,None,:,None]
10 loops, best of 3: 80.9 ms per loop

In [17]: %timeit np.einsum('ijk,lj->lijk',A,B)
10 loops, best of 3: 109 ms per loop

# Original soln
In [18]: %%timeit
    ...: d = []
    ...: for row in B:
    ...:     d.append(np.expand_dims(row[None,:,None]*A, axis=0))
    ...: 
    ...: result = np.vstack(d)
10 loops, best of 3: 130 ms per loop

Кредитное плечо multi-core

Мы могли бы использовать многоядерные возможности numexpr, которые подходят для arithmetic operations и large data, и таким образом получить некоторое повышение производительности здесь.Давайте рассмотрим время -

In [42]: import numexpr as ne

In [43]: B4D = B[:,None,:,None] # this is virtually free

In [44]: %timeit ne.evaluate('A*B4D')
10 loops, best of 3: 64.6 ms per loop

В одной строке: ne.evaluate('A*B4D',{'A':A,'B4D' :B[:,None,:,None]}).

Related post о том, как управлять многоядерными функциями.