Векторизованное точечное произведение намного медленнее, чем версия для цикла

Question

Ich хотел оптимизировать критическую часть моей программы, где большая часть работы выполняется. Эта работа состоит в вычислении многих точечных продуктов.

В настоящее время я делаю что-то вроде этого:

W = [np.random.rand(n_output, n_input) for _ in range(n_child)]
x = np.random.rand(n_input, batch_size)

for _ in range(n_times):
    for i in range(n_child):
        result[i] = np.dot(W[i],x)
    # work with result

Я хотел оптимизировать эту часть, избавившись от петли for i in range(n_child). Поэтому я сделал следующее:

W = np.random.rand(n_child, n_output, n_input)
x = np.random.rand(n_input, batch_size)

for _ in range(n_times):
    results = np.dot(W,x)
    # work with result

Однако это оказалось намного медленнее. Я что-то пропустил? В приведенном ниже примере мой второй подход примерно в 30 раз медленнее. Как это возможно?

Вот полностью рабочий пример:

import numpy as np
import time

n_child = 32
n_input = 1000
n_output = 20
batch_size = 64

W = np.random.rand(n_child, n_output, n_input)
x = np.random.rand(n_input, batch_size)

n_times = 1000

t0 = time.time()
for _ in range(n_times):
    np.dot(W,x)
t_a = time.time() - t0
print(t_a) # takes about 60 seconds on my machine

W = [np.random.rand(n_output, n_input) for _ in range(n_child)]

t0 = time.time()
for _ in range(n_times):
    for i in range(n_child):
        np.dot(W[i],x)
t_b = time.time() - t0
print(t_b) # takes about 3 seconds on my machine