Czym jest PyTorch Lightning i czym się różni od zwykłego PyTorcha?

PyTorch Lightning jest niczym innym, jak pythonowym frameworkiem do uczenia głębokiego. Jego nazwa może słusznie kojarzyć się z innym frameworkiem, PyTorchem. NIe jest to przypadkowa zbieżność, ponieważ Lightning opiera się na PyTorchu i skupia się głównie na organizacji kodu napisanego we wspomnianym frameworku. Z tego powodu dobrze sprawdza się przy bardziej skomplikowanych procesach uczenia modeli. Przez to wszystko Lightning jest też często określany wrapperem do PyTorcha.

Stworzenie nowego projektu - od czego zacząć?

Stworzenie nowego projektu z Lightningiem nie różni się zbytnio od stworzenia jakiegokolwiek innego projektu pisanego w Pythonie. Najważniejsze jest to, aby przed instalacją projektu stworzyć odrębne wirtualne środowisko. Przy instalacji PyTorcha przez pipa warto również dodać adres strony samego PyTorcha, aby pip instalował framework bezpośrednio stamtąd, a nie z PyPi, razem z całym wianuszkiem zależności.

Jak najlepiej ustrukturyzować nasz projekt? To zależy głównie od tego, w jaki sposób będziemy chcieli uczyć i optymalizować nasz model. Przedstawiony poniżej przykład jest najbardziej uniwersalną wersją projektu, którą łatwo można rozbudowywać w zależności od bieżących potrzeb.

Można cały szablon podzielić na trzy części

• katalog z danymi, logi oraz notatniki : rzeczy, które normalnie znajdują się w każdym projekcie MLowym, lecz na które Lightning nie ma większego wpływu

• katalog configs : tutaj znajdują się pliki konfiguracyjne projektu. Konfigurację projektu można zapisać w jednym wielkim pliku .yaml bądź porozdzielać go na klika mniejszych plików, co jest czytelniejsze przy większych projektach

• katalog src : w nim znajdują się pliki źródłowe projektu. Można go podzielić na kilka podstawowych modułów: moduł do trenowania modelu (modules), moduł do zarządzania danymi (datamodules), moduł z modelami (models) a także dodatkowe moduły z kodem własnych metryk czy funkcjami pomocniczymi (np. w przypadku przetwarzania języka naturalnego mogą to być funkcje pomocnicze do tokenizacji)

Ta ostatnia część jest najważniejsza, ponieważ przy niej najlepiej widać zalety Lightninga.

 

Organizacja procesu uczenia

Powyższa ilustracja została opracowana na podstawie animacji w dokumentacji Lightninga, którą możecie zobaczyć tutaj:

https://lightning.ai/docs/pytorch/stable/starter/introduction.html.

Po lewej znajduje się kod z PyTorcha, natomiast po prawej z Lightninga. Kolorami zostały zaznaczone bloki kodu odpowiedzialne za te same funkcjonalności. Pisanie całego procesu uczenia w postaci jednego długiego skryptu sprawiało, że w trakcie eksperymentów jakiekolwiek zmiany np. w strukturze modelu, czy trening na innych danych wymagał zmian całego skryptu oraz jego weryfikacji. Przy każdym projekcie powtarzały się te same czynności, np. trzeba było się upewnić, czy dany tensor na pewno został załadowany do odpowiedniego urządzenia przez wywołanie funkcji .to(device) . Jednym z celów Lightninga było ustrukturyzowanie tego kodu, aby był bardziej elastyczny i mógł być dostosowywany do bardziej skomplikowanych procesów uczenia. W tym celu zdefiniowano coś takiego jak system.

Warto wspomnieć, że LightningModule to nie jest to samo, co nn.Module z klasycznego PyTorcha. nn.Module definiuje pojedynczy model, natomiast LightningModule składa się z kilku modeli i bardziej określa środowisko, w którym znajdują się te modele oraz jakie relacje zachodzą między nimi.

Zarządzanie zbiorami danych

PyTorch domyślnie definiuje dwie klasy do zarządzania danymi - Dataset oraz DataLoader. Datasety służą do przechowywania danych wraz z ich etykietami, natomiast Dataloadery umożliwiają iterowanie po danych z Datasetów. Przed załadowaniem danych do DataLoaderów, należało je podzielić na zbiór treningowy, zbiór walidacyjny oraz zbiór testowy, a także przeprowadzić na nich odpowiednie transformacje. Lightning upraszcza cały proces zarządzania danymi dzięki klasie LightningDataModule.

Z podstawowych funkcji tej klasy można wyróżnić funkcję setup, w której można przeprowadzić transformacje na danych, a także podzielić je na odpowiednie zbiory. Istotnymi funkcjami są również train_dataloader, val_dataloader oraz test_dataloader, które, jak sama nazwa sugeruje, zwracają odpowiednie dataloadery.

Źródło: https://www.assemblyai.com/blog/pytorch-lightning-for-dummies/

Konfiguracja projektu przez CLI

Zmienianie parametrów procesów uczenia (w tym kontekście architektury modelu, zbioru danych czy hiperparametrów) przez konsolę jest często stosowane w celu uniknięcia potrzeby hardkodowania ich bezpośrednio w skryptach. Można w tym celu wykorzystać wbudowaną w Pythona klasę Argparse, która umożliwia przetwarzanie parametrów skryptu.

Lightning w tym obszarze również posiada swoją własną klasę LightningCLI. Pozwala ona na wywoływanie skryptu z różnymi plikami konfiguracyjnymi (zapisanymi we wspomnianym katalogu configs), a także z domyślnymi podkomendami, służącymi do trenowania, walidacji czy testowania modelu.

Powyżej znajduje się przykład wywołania skryptu z konfiguracją oraz z nadpisaniem parametru uczenia. Główną zaletą LightningCLI jest to, że użytkownik nie musi sam implementować własnego CLI, tylko może skupić się od razu na samych parametrach uczenia.

Śledzenie metryk za pomocą własnych bądź zewnętrznych loggerów

Przy procesie uczenia i walidacji modelu ważne jest śledzenie metryk modelu (np. dokładności) oraz zmian wartości funkcji straty w celu określenia, jaka jest jego jakość. Lightning pozwala na zapisywanie tych rzeczy za pomocą klasy Logger. W zależności od potrzeb użytkownika, można dynamicznie decydować o tym, jaki logger ma być używany przy uczeniu. Logi mogą być zapisywane do pliku .csv, wypisywane w konsoli bądź przekazywane do zewnętrznych narzędzi służących do logowania. Najczęściej jest to TensorBoard lub Wandb (Weights and Biases). Poniżej przedstawiony został przykład kodu wykorzystującego Wandb

Z perspektywy aplikacji wyniki zebrane przez loggera mogą wyglądać tak

Źródło: https://docs.wandb.ai/quickstart

Podsumowując, PyTorch Lightning powstał głównie z myślą o ustrukturyzowaniu i organizacji kodu, a co za tym idzie, większej elastyczności w rozwijaniu projektów MLowych. Mam nadzieję, że spodobała wam się prezentacja i jeżeli macie jakieś pytania bądź uwagi, zapraszam do ich zgłaszania. Wszelki feedback mile widziany.

Katarzyna Matuszek