Статья Автор: Деникина Н.В., Деникин А.В.

Векторизация

Почему цикл - это дорого

Четыре заказа - цикл незаметен. Но сервис доставки пересчитывает время для десятков тысяч заказов в секунду, а при обучении прогоняет миллионы строк истории десятки тысяч раз. Разница между «минута» и «неделя» появляется именно здесь.

Умножение матриц (300, 300) на (300, 100) - девять миллионов умножений.

В цикле Python каждое обходится в десятки машинных операций: проверить типы, создать объект, положить в переменную. Само умножение занимает ничтожную долю этого времени.

numpy отдаёт ту же работу библиотеке на C, которая считает блоками и укладывает данные в кэш процессора.

Векторизация - запись вычисления целиком над массивами, без цикла по отдельным числам.

Три вложенных цикла перебирают 300 примеров, 100 нейронов и 300 входов. Сколько умножений выполнится?

Проверить ответ

Три версии одного и того же

import numpy as np

rng = np.random.default_rng(0)
X = rng.normal(size=(300, 300))      # 300 примеров, 300 признаков
W = rng.normal(size=(300, 100))      # 300 входов, 100 нейронов

Y_slow = np.zeros((300, 100))
for i in range(300):                 # три вложенных цикла
    for j in range(100):
        s = 0.0
        for k in range(300):
            s = s + X[i, k] * W[k, j]
        Y_slow[i, j] = s

Y_mid = np.zeros((300, 100))
for i in range(300):                 # цикл только по примерам
    Y_mid[i] = X[i] @ W

Y_fast = X @ W                       # ни одного цикла

print(np.allclose(Y_slow, Y_fast))
print(np.allclose(Y_mid, Y_fast))

Вывод

True
True

Сколько это стоит по времени

import numpy as np
import time

rng = np.random.default_rng(0)
X = rng.normal(size=(300, 300))
W = rng.normal(size=(300, 100))

t0 = time.perf_counter()
Y_slow = np.zeros((300, 100))
for i in range(300):
    for j in range(100):
        s = 0.0
        for k in range(300):
            s = s + X[i, k] * W[k, j]
        Y_slow[i, j] = s
t1 = time.perf_counter()

t2 = time.perf_counter()
Y_fast = X @ W
t3 = time.perf_counter()

print("три цикла, с:", round(t1 - t0, 2))
print("без циклов, с:", round(t3 - t2, 5))
print("во сколько раз быстрее:", round((t1 - t0) / (t3 - t2)))

Вывод на нашей машине

три цикла, с: 3.3
без циклов, с: 0.00093
во сколько раз быстрее: 3548

Числа зависят от машины. Порядок останется тот же: секунды против долей миллисекунды.

Цикл по примерам - почти всегда лишний.

Как убирать циклы

Правило простое: если внутри цикла одна и та же операция над разными элементами, в numpy она уже есть.

Сумма - sum. Максимум с нулём - maximum. Сумма произведений - @. Среднее по столбцам - mean с axis.

Цикл остаётся нужен там, где шаги зависят друг от друга: следующий использует результат предыдущего. Обучение сети - как раз такой случай, цикл по шагам обучения убрать нельзя. А вот цикл по заказам внутри одного шага убирается всегда.

Отсюда рабочая привычка: увидел вложенные циклы по данным - сначала спроси, есть ли это одной операцией, и только потом пиши.

Убрав цикл, обязательно сверь результат через allclose. Векторная версия часто считает не совсем то же самое, и без сверки ты этого не заметишь.

И не оставляй np.append внутри цикла: он копирует весь массив на каждом шаге, и цикл становится медленнее, а не быстрее.

Печать