README.md

    Loss Landscape Analysis

    Loss Landscape Analysis (LLA) это библиотека для PyTorch, которая позволяет визуализировать и анализировать поверхности лосс нейросетей. Её основным отличием от существующих библиотек является наличие разнообразных опций для построения, которые включают в себя возможность выбора векторов-осей, уравнений апдейта весов, применения различных методов нормализации и т.д. LLA также предоставляет инструменты для анализа поверхностей лосс и гессиана нейросетей.

    Disclaimer: LLA находится в стадии разработки, из-за чего существующие функции могут меняться и новые функции могут быть добавлены в будущем. Рекомендуется использовать наиболее новую версию библиотеки.

    Требования

    • Python>=3.8.10
    • numpy>=1.22.4
    • torch>=1.13.0
    • torchvision>=0.14.0
    • matplotlib>=3.7.5
    • (опционально) jupyterlab>=3.2.9

    Запуск с Docker

    Для пользователей Docker и Docker Compose предоставляется удобный способ работы с LLA. Требуемые версии начинаются с Docker 27.3.1 и Docker Compose 2.29.7.

    Клонируйте этот репозиторий и запустите контейнер

    • git clone https://gitflic.ru/project/kryptodpi/loss-landscape-analysis.git
    • (опционально) измените порт в докер docker-compose.yml
    • docker compose up -d {service_name}

    В качестве {service_name} можно использовать следующие аргументы, которые влияют на использование GPU и Jupyter Lab

    • lla-cpu
    • lla-cpu-jup
    • lla-gpu
    • lla-gpu-jup

    В контейнере можно работать двумя способами

    • через терминал: docker exec -it lla_container /bin/bash
    • через Jupyter Lab: http:{server}:{your_port}, password land

    Примеры

    В репозиторий входит набор примеров examples.sh, которые иллюстрируют возможности построения поверхностей лосс и анализа гессиана с LLA. Данный шел скрипт сам запустит нужные скрипты и сохранит результаты в src/example_viz_results (изображения) и src/example_analysis_results (численные результаты). Обратите внимание, что эти директории очищаются при запуске скрипта, поэтому их не стоит использовать для хранения своих результатов. Используя Docker можно запустить данный скрипт командой

    docker exec -it lla-container /bin/bash examples-full.sh
    

    Выполнение данного скрипта может занять довольно много времени на ЦПУ, поэтому в случае отсутствия GPU можно запустить более короткую версию

    docker exec -it lla-container /bin/bash examples-fast.sh
    

    Конкретно для работы с LLA в репозиторий входят два скрипта lla_train.py и lla_eval.py для тренировки и инференса нейросетей. Эти скрипты также можно использовать для ознакомления с возможностями библиотеки. Ниже приведены примеры запуска с разными аргументами

    # use predefined plot settings with "all_modes", hessian axes, HESD, hessian criteria
    python3 lla_eval.py --cuda --seed 42 -lc 10 --weights ./example_weights/lenet_example.pth --all_modes --axes hessian --hessian -hc --name eval_example -vd viz_results -rd analysis_results
    # train for 2 epoch, visualize every 2 epochs, 9 plots per figure, random axes
    python3 lla_train.py --cuda --seed 42 --axes random --losscap 100 -np 1 -npp 1 --name random_example -vd viz_results -ld train_logs
    # train for 4 epochs, visualize every 2 epochs, use adam axes and adam update equation, filter L2 normalization
    python3 lla_train.py --cuda --seed 42 -e 4 -vr 2 --axes adam --mode adameq --norm filter --name adam_example -vd viz_results -ld train_logs
    

    Также есть отдельный скрипт esd_example.py для демонстрации возможностей анализа гессиана без построения поверхностей лосс. Ниже приведён пример запуска

    python3 esd_example.py --cuda --seed 42 --hessian -hc --name hessian_example -vd viz_results -rd analysis_results
    

    Также предоставлены Jupyter notebooks в директории src с примерами тренировки и инференса нейросетей, а также пример того, как LLA может быть импортирована в существующий проект. В этих файлах также можно увидеть примеры визуализаций, сделанные с LLA.

    Функции LLA

    Основной функцией LLA является визуализация 3D поверхностей лосс-функций и анализ спектра гессиана (HESD) нейросетей. LLA позволяет выбирать тип осей для построения поверхностей, уравнение апдейта весов, степень изменения весов, метод нормализации векторов-направлений. LLA также позволяет заморозить определенные слои нейросети и включает ряд “quality of life” опций. Все опции и аргументы функций перечислены в py_arguments-readme.md и viz_lla_esd-readme.md. Особое внимание нужно уделить двум аргументам, которые отличают LLA от существующих библиотек

    axes позволяет выбрать между случайными осями (random), осями Адама (adam) и осями гессиана (hessian) при построении поверхностей лосс. Случайные оси являются наиболее распространенным методом построения и реализованы во всех существующих библиотеках. Однако оси Адама, которые используют вектора-моменты оптимизатора Адам, и оси гессиана, которые используют собственные вектора, соответствующие максимальным собственным значениям гессиана, не встречаются в других библиотеках.

    mode позволяет выбрать уравнение апдейта весов между сложением (add) и уравнением оптимизатора Адам (adameq). Метод сложения является наиболее популярным и реализован в других библиотеках, а метод Адама предложен для апдейта весов при построении лосс впервые в данной библиотеке. Использование adameq также требует оси Адама.

    Кроме построения HESD спектров LLA позволяет рассчитать значения критериев анализа гессиана, которые предложены в [1].

    1. Импорт LLA в существующий проект

    Для использования LLA в своём проекте необходимо скопировать src/src_lla в желаемую директорию. Основной функционал LLA собран в функциях viz_lla и viz_esd, и пример их использования для анализа нейросети LeNet на датасете MNIST можно найти в import_example.ipynb.

    1.1. Визуализация поверхностей лосс с viz_lla

    Для использования viz_lla импортируйте его из библиотеки

    from src_lla import viz_lla, metrics
    

    viz_lla требует несколько объектов для работы:

    • model – объект нейросетевой модели в torch (наследующий от nn.Module);
    • metric – объект для расчета лосс, подробнее ниже;
    • (опционально) optimizer – объект оптимизатора torch.optim или подобный.

    Metric это объект LLA, который содержит всю необходимую информацию для расчета лосс. Он также использует батч данных для расчетов, который например можно взять из torch объекта DataLoader

    x_plot, y_plot = iter(data_loader).__next__()
    

    При работе с типовыми моделями объект metric требует только указания типа лосс-функции

    criterion = torch.nn.CrossEntropyLoss() # specify your loss function
    metric = metrics.Loss(criterion, x_plot, y_plot, device) # device - 'cuda' or 'cpu'
    

    Примеры вызовов viz_lla

    viz_lla(model,metric,device=device,normalization='filter',axes='random',to_save=True,to_viz=True)
    viz_lla(model,metric,device=device,axes='hessian',viz_dev=True,eval_hessian=True,to_save=True,to_viz=True)
    

    1.2. Анализ гессиана с viz_esd

    В то время как viz_lla позволяет использовать некоторые инструменты анализа гессиана, например, расчет и визуализацию HESD, основные функции анализа гессиана собраны в функцию viz_esd, которая может быть импортирована из библиотеки

    from src_lla import viz_esd
    

    Эта функция возвращает лист из пяти элементов, в которые входят собственные значения, собственные вектора, след, и критерии re и Khn (см. [1]). Ниже приведены примеры вызова viz_esd (полный список аргументов находится в viz_lla_esd-readme.md)

    # top 2 eigenvalues and eigenvectors
    eigvals, eigvects, _, _, _ = viz_esd(model, metric, esd=False, eigs=True, top_n=2)
    # eigenvalues, eigenvectors, trace, HESD and its criteria
    eigvals, eigvects, trace, re, Khn = viz_esd(model,metric,esd=True,eigs=True,trace=True,calc_crit=True,to_save=True,to_viz=True)
    # only HESD
    viz_esd(model, metric, esd=True, eigs=False, to_save=True, to_viz=True)
    

    2. Использование loaders для нестандартных нейросетей

    Если ваш проект использует нестандартные лосс-функции, даталоадеры, или инференс модели имеет вид более сложный, чем pred=model(x), то рекомендуется использование loaders. Они позволяют осуществить визуализацию поверхностей лосс практически для любого пайплайна, который можно реализовать в PyTorch. Создать loader для своего проекта можно по аналогии с loader_template.py и использовать его в lla_train.py или lla_eval.py импортируя по схеме, указанной ниже.

    2.1. Элементы в loader.py

    loader нужен для того, чтобы записать PyTorch пайплайн любой сложности в стандартизованном формате, что позволит использовать его другим объектам LLA. Это осуществляется путём создания нескольких стандартных элементов

    • ModelInit – функция, которая отвечает за инициализацию нейросетевой модели и возвращает её объект. От юзера требуется модифицировать строку model = _init_your_model(N_CLASSES) и добавить всё необходимое для инициализации модели и (опционально) загрузки весов.

    • CustomLoader – функция, которая отвечает за загрузку данных и возвращает объект DataLoader. От юзера требуется модифицировать строку train_data = _load_your_data(data_path) и добавить всё необходимое для загрузки своих данных и (опционально) препроцессинга. Также необходимо указать batch_size для объекта DataLoader. Важно, что DataLoader должен возвращать tuple inputs, labels. Если объект dataset (и, следовательно, dataloader) имеет другой формат выхода, то нужно указать wrap=True и в методе __next__ класса LoaderWrapper (см. ниже) указать, как выход вашего dataloader может быть сконвертирован в inputs, labels.

    • LoaderWrapper – класс-обёртка для dataloader, которая позволяет изменить формат его выхода. От юзера требуется модифицировать строку inputs, labels = output[0] в методе __next__ таким образом, чтобы выходом стали inputs, labels.

    • CustomLoss – класс, который отвечает за инференс model и вычисление значений лосс-функции. В __init__ нужно указать x (inputs из dataloader), y (labels из dataloader) и device, опционально юзер может добавить любые параметры, которые могут понадобиться методу __call__. От юзера требуется указать, как модель вычисляет логиты изменив строки pred = model_wrapper.forward(inputs.to(self.device)) и pred = model(inputs.to(self.device)). Также строка loss = self.loss_fn(pred, targets.to(self.device)) должна быть модифицирована путём добавления желаемой лосс-функции. Остальной код модифицировать не стоит*.

    *техническая заметка: CustomLoss возвращает loss как loss.item() при use_wrapper=True, что нужно для построения поверхностей лосс. Однако, если указать use_wrapper=False и return_pred=True, то CustomLoss будет вычислять loss не «отцепляя» её от вычислительного графа torch, что позволит делать loss.backward() и использовать CustomLoss при тренировке моделей. Это позволит обучать модель и строить поверхности лосс практически не модифицируя код. Также можно использовать аргумент batch для подачи тренировочных данных, см. пример в lla_train.py.

    2.2 Импорт loaders в другие скрипты

    lla_train.py и lla_eval.py написаны так, что могут работать с любым loader, который удовлетворяет условиям, описанным в разделе 2.1. Примеры loaders предоставлены для mlp, LeNet и ResNet (см. скрипты в src/src_lla/loaders). После того, как loader импортирован в проект, все объекты могут быть инициализированы следующим образом

    from src_lla.loaders.{your_loader} import *
    
    # use standard loader elements to initialize all required components
    dataloader = CustomLoader()
    model = ModelInit(device=device)
    metric = CustomLoss(x_plot, y_plot, device)
    

    3. Известные особенности

    Первому вызову любой функции LLA требующей анализа гессиана сопутствует Warning. Он предупреждает о возможной утечке памяти, но может быть проигнорирован, так как библиотека позволяет избежать утечки памяти при использовании viz_lla и viz_esd. Однако, таких гарантий нет если юзер будет вызывать отдельные компоненты функций LLA, см. описание метода hessian_calc.reset() в src_lla/hessian.

    Цитирование LLA

    Если вы используете LLA в своей работе, пожалуйста, процитируйте данную статью:

    @misc{lla,
    title={Investigating generalization capabilities of neural networks by means of loss landscapes and Hessian analysis}, 
    author={Nikita Gabdullin},
    year={2024},
    url={https://arxiv.org/abs/2412.10146}}
    

    [1] N. Gabdullin, “Investigating generalization capabilities of neural networks by means of loss landscapes and hessian analysis,” 2024. [Online]. Available: https: //arxiv.org/abs/2412.10146

    License

    Copyright 2024 Kryptonite

    Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

    Конвейеры
    0 успешных
    0 с ошибкой