PyTorch vs TensorFlow: malý experiment

Na začátku stály čtyři praktické otázky: vyplatí se používat TensorFlow, stačil by PyTorch, jaké jsou reálné výhody a nevýhody a jak moc se rozdíly projeví v běžném projektu. Abychom ale zodpověděli na otázky bude zapotřebí i malý experiment.

Nastavení experimentu

Experiment byl navržen jako srovnání dvou frameworků na stejné úloze image recognition. Použit byl dataset CIFAR-10, který obsahuje 60 000 obrázků o velikosti 32 × 32 pixelů rozdělených do deseti tříd. Testované architektury byly ResNet50, VGG19 a MobileNet. Trénink probíhal 5 epoch s batch size 32 na hardwaru Intel i7-11370 a NVIDIA GeForce RTX 3060.

Postup implementace

V obou frameworkcích měl experiment stejnou logiku:

1. Načtení dat (train/test) z CIFAR10 datasetu
2. Inicializace modelu (ResNet50/VGG19/ MobileNet)
3. Predikce před tréninkem – slouží jako baseline
4. Fine-tuning modelu na datasetu
5. Predikce po tréninku

Co se během testu sledovalo

Samotná přesnost po tréninku byla jen jedna část obrazu. Smyslem testu bylo sledovat také to, jak se frameworky chovají během výpočtu. Sledovaly se tyto parametry:

• Přesnost modelu před a po fine-tuningu
• Celkový čas tréninku
• Vytížení GPU během běhu a chování frameworku mezi epochami
• Způsob práce s GPU pamětí
• Vývojářská zkušenost: instalace, čitelnost kódu a dokumentace

Výsledky experiment

Přesnost modelu před a po fine-tuningu

Celkový čas tréninku

Vytížení GPU během běhu a chování frameworku mezi epochami

Na těchto grafech je vidět, že Tensorflow má znatelně vyšší režii kolem epoch. Často za tím stojí procesy, které probíhají na konci epochy, jako třeba validace, přenastavování vah, callback, …

U krátkých běhů se tyto režie projevují vice, u dlouhých tréninků (desítky až stovky epoch) se fixní režie může „rozpustit“ a rozdíl se zmenší.

Způsob práce s GPU pamětí

Zde jde vidět výrazný rozdíl ve strategii obou framework. Pytorch je předvídatelnější při debugování, TensorFlow má tendenci si alokovat paměť po blocích.

Vývojářská zkušenost: instalace, čitelnost kódu a dokumentace