Narzut GPU: dlaczego nasz model MNIST trenuje się szybciej na CPU
Uczenie maszynowe i GPU wydają się nierozłączne — każdy tutorial, każdy dostawca chmury, każdy przewodnik „jak zacząć” kieruje Cię ku instancjom z GPU. Intuicja jest prosta: dane wchodzą, tysiące rdzeni przetwarza je równolegle, trening idzie szybko. Ale zakłada to, że wszystkie te rdzenie są rzeczywiście używane — a w przypadku małych obciążeń nie są.
Wytrenowaliśmy nasz model MNIST na lokalnej stacji roboczej z GPU Quadro RTX 5000 i procesorem Intel Core i9-10885H, a następnie tylko na CPU. CPU był szybszy — 4,4 s wobec 6,5 s na epokę. Ten artykuł wyjaśni dlaczego i pokaże, jak zbiliśmy trening do 0,14 s, redukując narzut zamiast wymieniać sprzęt.
Zacznijmy od tego, co dzieje się podczas treningu.
Rozkładamy jeden krok treningu
Każdy krok treningu przebiega w czterech fazach — przejście w przód, strata, przejście w tył, aktualizacja wag:
# Przejście w przód
z1 = xb @ w1 + b1 # mnożenie macierzy + bias (warstwa ukryta)
a1 = np.maximum(0, z1) # aktywacja ReLU
z2 = a1 @ w2 + b2 # mnożenie macierzy + bias (warstwa wyjściowa)
exp_z = np.exp(z2 - z2.max(axis=1, keepdims=True))
probs = exp_z / exp_z.sum(axis=1, keepdims=True) # softmax → prawdopodobieństwa
# Gradient straty
dz2 = probs.copy()
dz2[np.arange(bs), yb] -= 1 # jak bardzo przewidywania się pomyliły
dz2 /= bs
# Przejście w tył — obliczamy gradienty
dw2 = a1.T @ dz2 # gradient dla W2
da1 = dz2 @ w2.T # gradient płynący wstecz
dz1 = da1 * (z1 > 0) # gradient ReLU
dw1 = xb.T @ dz1 # gradient dla W1
# Aktualizujemy wagi
w1 -= lr * dw1
b1 -= lr * dz1.sum(axis=0)
w2 -= lr * dw2
b2 -= lr * dz2.sum(axis=0)Gdzie ucieka czas? Żeby to sprawdzić, owinęliśmy każdą operację w time.perf_counter():
t = time.perf_counter()
z1 = xb @ w1 + b1
timings["fwd: X @ W1 + b1"] += time.perf_counter() - t
t = time.perf_counter()
a1 = np.maximum(0, z1)
timings["fwd: ReLU"] += time.perf_counter() - t
# ... i tak dalej dla każdej operacjiPrzy uruchamianiu na GPU musieliśmy dodać
tf.test.experimental.sync_devices()między operacjami — bez tego operacje GPU kolejkują się asynchronicznie i pomiar rejestruje tylko zlecenie, a nie faktyczne wykonanie.
Oto wyniki dla jednego kroku treningu (batch_size=32, średnia z 1000 kroków, czas w mikrosekundach μs — milionowych częściach sekundy, mniej znaczy lepiej):
| Operacja | CPU (NumPy) | GPU (TF) |
|---|---|---|
| fwd: X @ W1 + b1 (obrazy × wagi warstwy ukrytej + bias) | 162 μs | 541 μs |
| fwd: ReLU (zerowanie wartości ujemnych) | 9 μs | 173 μs |
| fwd: a1 @ W2 + b2 (warstwa ukryta × wagi wyjścia + bias) | 19 μs | 478 μs |
| fwd: softmax + strata (prawdopodobieństwa + błąd) | 30 μs | 780 μs |
| bwd: gradienty (ile każda waga wniosła do błędu) | 188 μs | 1271 μs |
| update: W -= lr*dW (korekta wag, by zmniejszyć błąd) | 345 μs | 1204 μs |
| RAZEM | 751 μs | 4447 μs |
Pierwsze, co się rzuca w oczy: CPU jest szybszy w każdej pojedynczej operacji — 751 μs łącznie wobec 4447 μs. Operacja, która powinna zyskać na zrównoleglaniu na GPU najwięcej, to mnożenie macierzy X @ W1 — to tysiące niezależnych iloczynów skalarnych, które mogłyby policzyć się jednocześnie na tysiącach rdzeni, dokładnie taka praca, do jakiej GPU zostały zaprojektowane.
A jednak w tabeli powyżej wciąż jest wolniejsza na GPU: 162 μs na CPU wobec 541 μs na GPU. CPU jest tu szybki, bo NumPy woła bezpośrednio do BLAS (Basic Linear Algebra Subprograms) — mocno zoptymalizowanych procedur w C/Fortranie, które korzystają ze specyficznych dla CPU instrukcji SIMD (AVX2 przetwarza 8 liczb zmiennoprzecinkowych na cykl, FMA scala mnożenie z dodawaniem w jedną instrukcję). GPU jest wolniejsze nie dlatego, że jego arytmetyka jest wolniejsza, ale ponieważ każda operacja płaci narzut CUDA (uruchomienie kernela, synchronizacja pamięci, przełączanie kontekstu).
Gdy ten narzut zabiera więcej czasu niż same obliczenia, GPU wychodzi wolniejsze, mimo że liczy szybciej.
Spójrzmy na xb @ w1 + b1 i a1 @ w2 + b2 — mnożenia macierzy z każdej warstwy — i na to, dlaczego GPU są projektowane, by robić je szybko. Oto przejście w przód z naszej implementacji w NumPy:
class HiddenLayer:
def forward(self, x):
self.z = self.W @ x + self.b # mnożenie macierzy + bias
self.out = np.maximum(0, self.z) # aktywacja ReLU
return self.out
class OutputLayer:
def forward(self, x):
self.z = self.W @ x + self.b # mnożenie macierzy + bias
exp = np.exp(self.z - np.max(self.z))
self.probs = exp / np.sum(exp) # softmax → prawdopodobieństwa
return self.probsKażde @ to mnożenie macierzy. Dla jednego obrazu self.W @ x mnoży macierz wag przez piksele obrazu. Każdy ze 128 neuronów ma 784 wagi — po jednej na piksel wejściowy.
Każdy neuron liczy iloczyn skalarny swoich 784 wag z 784 wartościami wejściowymi (pikselami obrazu) — to 784 mnożenia i 783 dodawania na neuron. Wynikiem jest 128 aktywacji, pochodzących od 128 neuronów warstwy, z których każdy daje jedną wartość ze swojego iloczynu skalarnego. To 128 niezależnych iloczynów skalarnych na jeden obraz:
x (784,) W (128 × 784) result (128,)
┌─────────────┐ ┌──────────────────┐ ┌──────────────┐
│ p1 p2 … p784│ @ │ n1: w1 w2 … w784│ = │ a1 │
└─────────────┘ │ n2: w1 w2 … w784│ │ a2 │
│ ... │ │ ... │
│ n128:w1 w2 … w784│ │ a128 │
└──────────────────┘ └──────────────┘Po lewej — jeden obraz o 784 pikselach, w środku — 128 neuronów, po prawej — 128 aktywacji.
Ale frameworki takie jak Keras nie przetwarzają obrazów po jednym — układają cały batch w jedną macierz i mnożą wszystko naraz. Przy rozmiarze batcha 32 X ma 32 wiersze × 784 kolumny, a ta sama macierz wag daje 32 × 128 = 4096 aktywacji w jednej operacji:
X (32 × 784) W1 (784 × 128) result (32 × 128)
┌─────────────────┐ ┌──────────────────┐ ┌──────────────────┐
│ img1: p1 … p784│ │ n1 n2 … n128 │ │ img1: a1 … a128 │
│ img2: p1 … p784│ @ │ w w … w │ = │ img2: a1 … a128 │
│ ... │ │ ... ... ... │ │ ... │
│ img32: p1 … p784│ │ w w … w │ │ img32:a1 … a128 │
└─────────────────┘ └──────────────────┘ └──────────────────┘Po lewej — 32 obrazy po 784 piksele każdy, w środku — 784 wagi na neuron, po prawej — 32 × 128 aktywacji.
Wagi są wspólne — to tak, jakbyś miał 32 kopie tego samego modelu działające jednocześnie na 32 obrazach, tylko że robi się to jedną operacją, a nie 32 osobnymi. Jedno uruchomienie kernela, jedno odczytanie macierzy wag z pamięci i wszystkie 4096 wyników policzone równolegle.
Teraz ciekawe pytanie: ile z tych 4096 iloczynów skalarnych GPU potrafi policzyć jednocześnie? Ponieważ każdy z nich jest całkowicie niezależny — aktywacja neuronu 20 dla obrazu 5 nie zależy od żadnego innego wyniku — GPU może przypisać każdy do osobnego wątku.
RTX 5000 ma 48 multiprocesorów strumieniowych (SM), każdy z 64 rdzeniami CUDA — łącznie 3072 rdzenie. Każdy rdzeń pracuje z częstotliwością ~1,8 GHz i może wykonać mnożenie z dodawaniem na cykl, co daje teoretyczny szczyt: 3072 rdzeni × 1,8 mld cykli/s × 2 operacje/cykl = ~11 TFLOPS (bilionów operacji zmiennoprzecinkowych na sekundę).
Nasze mnożenie macierzy daje 4096 iloczynów skalarnych, każdy wymagający 784 mnożeń z dodawaniem: 4096 × 784 = ~3,2 miliona operacji łącznie. GPU potrafi wykonać 11 bilionów operacji na sekundę, a my prosimy tylko o 3,2 miliona — to 0,3 mikrosekundy faktycznych obliczeń, czyli zaledwie 0,00003% możliwości GPU. GPU kończy arytmetykę w 0,3 μs, a potem czeka ~540 μs na kolejne uruchomienie kernela. Pracuje 0,06% czasu.
Gdzie ten czas naprawdę idzie
Trenowanie sieci neuronowej na GPU to więcej niż mnożenie macierzy. Każdy krok treningu jest ciągiem operacji, a dla małych modeli faktyczna arytmetyka jest jego najmniejszą częścią:
1. Narzut uruchomienia kernela. Każda operacja GPU — mnożenie macierzy, funkcja aktywacji, obliczenie straty — jest kernelem, który CPU musi zaplanować i uruchomić na GPU. Każde uruchomienie ma stały narzut rzędu 5–15 mikrosekund. Dla dużego mnożenia macierzy trwającego milisekundy jest to pomijalne. Dla naszego maleńkiego mnożenia 32×784 × 784×128, kończącego się w mikrosekundach, narzut uruchomienia może przewyższyć same obliczenia. Samo nasze przejście w przód obejmuje mnożenie macierzy, dodanie biasu, ReLU, kolejne mnożenie macierzy, kolejny bias i softmax — co najmniej 6 uruchomień kerneli, jeszcze przed rozpoczęciem propagacji wstecznej.
2. Opóźnienie transferu pamięci. Dane muszą podróżować z pamięci CPU do pamięci GPU (a gradienty z powrotem). Ten transfer ma stałe opóźnienie — czas na przygotowanie transferu DMA, przejście przez szynę PCIe i zasygnalizowanie zakończenia. Przy rozmiarze batcha 32 i wejściach 784-wymiarowych przesyłamy około 100 KB na batch. Szyna PCIe potrafi przenieść 32 GB/s, więc surowy transfer to ~3 mikrosekundy — ale narzut przygotowania jest 10–20× większy.
3. Narzut Pythona i frameworka. Keras/TensorFlow dodają własną warstwę pośrednictwa. Każda operacja przechodzi przez Pythona, środowisko uruchomieniowe TensorFlow, kompilację XLA (kompilator optymalizujący graf obliczeń przy pierwszym uruchomieniu), alokację pamięci i synchronizację. Dla dużych operacji ten narzut jest niewidoczny. Dla małych to on jest wąskim gardłem.
Dobra wiadomość jest taka, że możemy dokładnie zmierzyć, ile czasu zabiera każdy typ narzutu.
Profilowanie narzutu CUDA
Nasze pomiary per operacja pokazały, że GPU zabiera więcej czasu, a wyliczyliśmy, że faktyczna arytmetyka zajmuje zaledwie 0,3 μs — GPU powinno więc przez większość czasu być bezczynne. Na co dokładnie spędza pozostałe 540 μs?
Żeby to sprawdzić, użyliśmy profilera NVIDII nsys (Nsight Systems), który przechwytuje każde wywołanie API CUDA:
nsys profile -o keras-gpu-profile python mnist-keras.py
nsys stats --force-export=true keras-gpu-profile.nsys-repCUDA to warstwa programowa NVIDII między Twoim kodem a sprzętem GPU. Gdy TensorFlow chce pomnożyć macierze, nie rozmawia z GPU bezpośrednio — woła funkcje CUDA w rodzaju „zaalokuj pamięć”, „skopiuj te dane”, „uruchom ten kernel”. Każde z tych wywołań przechodzi przez sterownik i ma własny narzut. Profiler nsys rejestruje każde z nich, więc możemy dokładnie zobaczyć, gdzie ucieka czas.
Tym razem, zamiast mierzyć poszczególne operacje, sprofilowaliśmy całą epokę model.fit() — wszystkie 1500 kroków. Epoka zajęła 6,5 s na GPU. Oto gdzie ten czas poszedł:
| Operacja | Typ narzutu | Wywołania | Czas łącznie |
|---|---|---|---|
cuCtxSetCurrent (przełączanie kontekstu) | uruchomienie kernela | 98 168 | 0,87 s |
cuEventRecord (zdarzenia pomiaru/synchronizacji) | uruchomienie kernela | 22 786 | 0,57 s |
cuMemcpyDtoHAsync (kopiowanie GPU→CPU) | transfer pamięci | 6144 | 0,56 s |
cuMemcpyHtoDAsync (kopiowanie CPU→GPU) | transfer pamięci | 5678 | 0,27 s |
cuGraphLaunch (wykonanie skompilowanego grafu) | uruchomienie kernela | 1875 | 0,16 s |
cuLaunchKernel | uruchomienie kernela | 1080 | 0,02 s |
| Razem narzut CUDA | ~2,7 s |
Żaden z tych wierszy CUDA nie jest samym mnożeniem macierzy — wszystkie to narzut zarządzania wokół niego. Faktyczne obliczenia (mnożenia macierzy, ReLU, softmax) dzieją się na GPU po tym, jak cuGraphLaunch zleci pracę, ale dla naszych maleńkich macierzy jest to tak szybkie, że nie pojawia się nawet jako istotna pozycja.
Składając wszystko razem, oto gdzie poszła cała epoka trwająca 6,5 s:
| Czas | |
|---|---|
| Narzut CUDA (tabela powyżej) | ~2,7 s |
| Narzut Pythona/frameworka (nierejestrowany przez nsys) | ~3,8 s |
| Faktyczna arytmetyka GPU | pomijalna |
| Cała epoka | 6,5 s |
W praktyce całe 6,5 s to narzut. Keras jest już bardzo dobrze zoptymalizowany — używa CUDA Graphs (cuGraphLaunch — 1875 wywołań, po jednym na batch), żeby wstępnie skompilować całe przejście w przód i w tył i odtwarzać je bez zlecania każdej operacji z osobna. Ale nawet z tą optymalizacją GPU nigdy nie dostaje szansy, by odrobić narzut szybszymi obliczeniami.
Dla porównania: Keras na CPU przechodzi tę samą epokę w 4,4 s — szybciej niż 6,5 s na GPU, bo narzutu CUDA nie ma wcale. XLA kompiluje do kodu natywnego, a CPU wykonuje arytmetykę wprost we własnej pamięci.
Warto podkreślić, że co do samej arytmetyki nie ma żadnej niespodzianki — GPU jest znacząco szybsze, zgodnie z oczekiwaniami. Zmierzyliśmy kluczową operację — jedno mnożenie macierzy (32, 784) @ (784, 128), jeden batch przez pierwszą warstwę:
| Czas na jedno mnożenie macierzy | |
|---|---|
| GPU | 109 μs |
| CPU | 272 μs |
GPU jest 2,5× szybsze — ale zauważ, że nawet te 109 μs zawierają już narzut CUDA na uruchomienie tej jednej operacji. Czysta arytmetyka to byłoby ~0,3 μs (jak wyliczyliśmy wcześniej z TFLOPS-ów); pozostałe ~108 μs to narzut tego jednego uruchomienia kernela. Na CPU całe 272 μs to arytmetyka — żadnej warstwy narzutu pomiędzy.
Pełny krok treningu obejmuje kilka mnożeń macierzy plus aktywacje, stratę, gradienty i całe zarządzanie CUDA wokół każdego z nich. Gdy mierzymy skompilowany train_step z całym narzutem, różnica znika kompletnie:
| Na krok | |
|---|---|
| GPU | 1,27 ms (~0,1 ms obliczeń + ~1,2 ms narzutu) |
| CPU | 1,30 ms (wszystko obliczenia, zero narzutu) |
GPU wykonuje arytmetykę szybciej, ale resztę czasu przeznacza na narzut, wychodząc w efekcie na podobną szybkość na krok. A w pełnym potoku model.fit() — z wczytywaniem danych, metrykami i callbackami — CPU wręcz wygrywa: 4,4 s wobec 6,5 s na epokę.
Ciekawa uwaga: jeśli uruchomisz nvidia-smi w trakcie treningu, możesz zobaczyć wykorzystanie GPU na poziomie 90–100%. Wygląda to, jakby GPU było w pełni zajęte — dlaczego więc jest wolniejsze od CPU?
Ponieważ nvidia-smi raportuje odsetek czasu, w którym na GPU działa co najmniej jeden kernel — a nie ile rdzeni jest aktywnych. Jeśli maleńkie kernele są uruchamiane jeden za drugim bez przestojów, pokaże ~100% wykorzystania, mimo że przytłaczająca większość rdzeni w każdej chwili stoi bezczynnie. Jak widzieliśmy wcześniej, nasz model używa zaledwie 0,00003% możliwości GPU.
Jak przyspieszyć jeden przebieg treningu
Skoro rozumiemy już narzut, jak go zredukować? Jest kilka kierunków ataku: zmniejszyć liczbę kroków (większy rozmiar batcha), wyeliminować narzut CUDA (działać na CPU), wyeliminować narzut frameworka (pominąć TensorFlow), skompilować cały krok w jedną operację (JIT w JAX) albo naiwnie przenieść się na GPU (CuPy — spoiler: jest gorzej). Wypróbowaliśmy wszystkie pięć.
1. Większy rozmiar batcha
Jeśli narzut jest płacony za każdy krok, oczywistą poprawką jest: mniej kroków. Większy rozmiar batcha oznacza więcej próbek na krok, więc ta sama epoka wymaga mniej kroków — a każdy krok płaci podatek narzutu tylko raz, niezależnie od rozmiaru batcha. Przy 48 000 próbek treningowych (60 tys. minus 20% na walidację) i rozmiarze batcha 32 jedna epoka to 48 000 / 32 = 1500 kroków, gdzie każdy krok to jedno pełne przejście w przód + w tył + aktualizacja na jednym batchu. Przy rozmiarze batcha 4096 to zaledwie 11 kroków.
Większe batche dają też GPU więcej pracy na krok — jak widzieliśmy wcześniej, rozmiar batcha 32 utrzymuje aktywne tylko ~3 z 48 SM, podczas gdy rozmiar 2048 zajmuje ~42. Większe batche pomagają więc na dwa sposoby: mniej kroków płacących narzut na epokę i więcej faktycznie pracujących rdzeni GPU na krok.
Nie można jednak zwiększać rozmiaru batcha bez opamiętania — ma on bezpośredni wpływ na dokładność modelu. Większe batche dają gładsze, ale rzadsze aktualizacje gradientu, co może prowadzić do gorszej generalizacji. Właściwy rozmiar batcha musisz wyeksperymentować dla swojego konkretnego modelu.
2. Keras na CPU — pominąć narzut CUDA
Jak widzieliśmy, samo wyłączenie GPU (tf.config.set_visible_devices([], 'GPU')) eliminuje 2,7 s narzutu CUDA. XLA kompiluje do kodu natywnego, a CPU wykonuje arytmetykę bezpośrednio.
3. Czysty NumPy — pominąć także narzut frameworka
Idąc dalej, możemy całkowicie zrezygnować z TensorFlow. Pętla treningowa w czystym NumPy woła bezpośrednio do procedur BLAS, bez narzutu zlecania frameworka między operacjami.
4. JIT w JAX — skompilować cały krok
Kompilator jit z JAX-a trasuje cały krok treningu i kompiluje go w jedną zoptymalizowaną funkcję natywną. Zamiast tego, by Python zlecał każdą operację po kolei, skompilowana funkcja wykonuje je wszystkie w jednym scalonym wywołaniu z niemal zerowym narzutem.
5. CuPy — a co, jeśli po prostu przeniesiemy NumPy na GPU?
Wypróbowaliśmy też podejście naiwne: zamienić import numpy as np na import cupy as cp i uruchomić ten sam kod na GPU. Ponieważ nasza implementacja przetwarza próbki po jednej w pętli Pythona, każda maleńka operacja (dodanie 128-elementowego wektora, softmax na 10 elementach) staje się osobnym uruchomieniem kernela GPU. Rezultat: 443 sekundy na 5 epok — ponad 6× wolniej niż NumPy na CPU. Samo przeniesienie kodu na GPU bez przemyślenia wzorca dostępu pogarsza sprawę, a nie poprawia.
Składamy to razem
Przetestowaliśmy podejścia 1–4 na różnych rozmiarach batcha (CuPy był zbyt wolny, by go włączyć). Każda komórka to czas jednej epoki:
| Rozmiar batcha | Keras GPU | Keras CPU | Czysty NumPy | JAX (CPU, JIT) |
|---|---|---|---|---|
| 32 | 6,26 s | 11,31 s | 2,71 s | 2,37 s |
| 128 | 2,74 s | 2,21 s | 4,70 s | 1,03 s |
| 512 | 2,64 s | 1,80 s | 1,30 s | 0,64 s |
| 1024 | 2,47 s | 1,43 s | 0,87 s | 0,63 s |
| 2048 | 2,32 s | 1,42 s | 0,84 s | 0,48 s |
| 4096 | 2,32 s | 1,39 s | 1,03 s | 0,14 s |
Czytanie wzdłuż wierszy pokazuje efekt zmiany podejścia (mniej narzutu na krok). Czytanie w dół kolumn pokazuje efekt większych batchy (mniej kroków na epokę). Oba efekty się kumulują.
Keras na GPU zatrzymuje się na ~2,3 s niezależnie od rozmiaru batcha — narzut CUDA ma stałą podłogę, której większe batche nie zlikwidują.
Keras na CPU schodzi z 11,3 s do 1,4 s — bez narzutu CUDA, ale narzut samego frameworka TensorFlow wyznacza własną podłogę.
Czysty NumPy osiąga 0,84 s przy rozmiarze batcha 2048 — żadnego frameworka, tylko wywołania BLAS. Przy 4096 znów zwalnia z powodu presji na pamięć.
JAX dochodzi do 0,14 s przy rozmiarze batcha 4096 — 45× szybciej niż Keras na GPU. Kompilacja JIT scala cały krok w jedno wywołanie natywne z niemal zerowym narzutem.
Wniosek: droga do szybszego treningu nie prowadzi przez większy sprzęt, ale przez mniejszy narzut. Każde podejście w tabeli zdejmuje jedną warstwę narzutu, a większe batche zmniejszają, ile razy płacisz za to, co zostało.
Zrównoleglanie przy przeszukiwaniu hiperparametrów
Powyższe podejścia przyspieszają jeden przebieg treningu. Ale gdy przeszukujesz hiperparametry — próbujesz 5 współczynników uczenia albo 4 architektury — każdy przebieg jest całkowicie niezależny. Nie dzielą wag, gradientów ani stanu.
vmap (zwektoryzowana mapa) z JAX-a potrafi trenować wszystkie 5 modeli jednocześnie w jednym przejściu w przód i w tył.
To coś innego niż przetwarzanie batchowe w Kerasie, które przepuszcza wiele obrazów przez jeden model.
vmap przepuszcza te same obrazy przez wiele modeli — każdy z własnymi wagami i własnym współczynnikiem uczenia — w jednej scalonej operacji. Pod spodem XLA kompiluje to w jeden kernel: jedno mnożenie macierzy obsługuje wszystkie 5 przejść w przód, drugie wszystkie 5 przejść w tył. Żadnej pętli Pythona, żadnego narzutu na model.
Kompromis: wszystkie modele muszą mieć tę samą architekturę (te same kształty parametrów), a pętlę treningową trzeba przepisać w JAX-ie. Oto pełny kod:
import jax
import jax.numpy as jnp
from jax import vmap, jit, grad, random
def init_params(key):
# Ten sam model co wcześniej: 784→128→10, losowe wagi
k1, k2 = random.split(key)
w1 = random.normal(k1, (784, 128)) * jnp.sqrt(2.0 / 784)
b1 = jnp.zeros(128)
w2 = random.normal(k2, (128, 10)) * jnp.sqrt(2.0 / 128)
b2 = jnp.zeros(10)
return (w1, b1, w2, b2)
def loss_fn(params, x, y):
# Przejście w przód + strata entropii krzyżowej — ta sama matematyka co w wersji NumPy
w1, b1, w2, b2 = params
h = jnp.maximum(0, x @ w1 + b1) # warstwa ukryta + ReLU
logits = h @ w2 + b2 # warstwa wyjściowa
log_probs = logits - jnp.log(jnp.sum(jnp.exp(logits), axis=-1, keepdims=True))
return -jnp.mean(log_probs[jnp.arange(y.shape[0]), y])
def sgd_step(params, x, y, lr):
# Jeden krok treningu: policz gradienty, zaktualizuj wagi
grads = grad(loss_fn)(params, x, y) # JAX automatycznie różnicuje loss_fn
return tuple(p - lr * g for p, g in zip(params, grads))
# 5 współczynników uczenia, 5 zestawów wag, trenowanych jednocześnie
lr_array = jnp.array([0.001, 0.01, 0.1, 1.0, 10.0])
batched_step = jit(vmap(sgd_step, in_axes=(0, None, None, 0)))
batched_params = batched_step(batched_params, x_batch, y_batch, lr_array)W prostszych przypadkach ten sam cel osiąga multiprocessing.Pool z Pythona — 5 osobnych procesów, każdy trenujący jedną konfigurację:
import multiprocessing as mp
def train_one_config(args):
# Trenuje jeden model z jednym współczynnikiem uczenia — działa we własnym procesie
lr, X_train, y_train = args
model = keras.Sequential([
keras.layers.Dense(128, activation="relu", input_shape=(784,)),
keras.layers.Dense(10, activation="softmax"),
])
model.compile(optimizer=keras.optimizers.SGD(learning_rate=lr),
loss="sparse_categorical_crossentropy")
history = model.fit(X_train, y_train, epochs=10, batch_size=32,
validation_split=0.2, verbose=0)
return lr, history.history
with mp.Pool(5) as pool:
results = pool.map(train_one_config,
[(lr, X_train, y_train) for lr in [0.001, 0.01, 0.1, 1.0, 10.0]])Co warto zapamiętać
- Dla małych modeli CPU jest szybszy niż GPU. Nasz model MNIST (784→128→10, ~101 tys. parametrów) trenował się w 4,4 s na epokę na CPU wobec 6,5 s na GPU. Arytmetyka GPU jest szybsza (109 μs wobec 272 μs na mnożenie macierzy), ale narzut CUDA — uruchomienia kerneli, kopiowania pamięci, przełączanie kontekstu — dodaje 2,7 s na epokę, których CPU po prostu nie ma.
- GPU jest prawie nieużywane. Nasze mnożenie macierzy wykorzystuje 0,00003% możliwości RTX 5000. Przy batch_size=32 aktywne są tylko ~3 z 48 SM. GPU kończy arytmetykę w 0,3 μs, a potem czeka ~540 μs na kolejne uruchomienie kernela.
- Wskaźnik wykorzystania w
nvidia-smiwprowadza w błąd. Raportuje czas, w którym działa jakikolwiek kernel, a nie obsadzenie rdzeni. Nasz model pokazuje ~100% wykorzystania, korzystając z 0,00003% mocy obliczeniowej GPU. - Mniej narzutu bije większy sprzęt. Przechodząc z Kerasa na GPU do JIT w JAX-ie na CPU przy batch_size=4096, zeszliśmy z 6,5 s do 0,14 s na epokę — 45-krotne przyspieszenie na tej samej maszynie, bez żadnego GPU.
- Nie sięgaj po GPU, dopóki Twoje obciążenie tego nie wymaga. Dla modeli z mniej niż ~500 tys. parametrów i rozmiarami batcha poniżej 256 szybki CPU będzie i tańszy, i szybszy.
Ten sam schemat objawia się głośno w produkcji w przypadku word2vec: każdy krok treningu dotyka zaledwie garstki wierszy dwóch macierzy zanurzeń i wykonuje rzędu 10⁴ mnożeń z dodawaniem — około tysiąc razy mniej obliczeń na przykład niż powyższe przejście w przód MNIST. Word2vec to obciążenie zdominowane przez dostęp do pamięci i losowe wyszukiwanie wierszy, a nie przez gęste mnożenie macierzy — czyli dokładnie taki kształt, który przegrywa na GPU i wygrywa na CPU. De facto standardowa biblioteka do word2vec, Gensim, idzie w to na całość: wewnętrzna pętla w Cythonie, bezblokadowe wątki CPU w duchu Hogwild!, żadnego GPU na żadnym etapie — i wciąż bije implementacje GPU pod względem czasu zegarowego przy trenowaniu tablicy wyszukiwania (V, d) z korpusu, ponad dekadę po tym, jak CUDA zaczęła dominować w uczeniu głębokim. Artykuł o word2vec szczegółowo wyjaśnia dlaczego.