Hessian-aware pruning and Optimal Neural Implant
В этом материале на примере статьи Hessian-aware pruning and Optimal Neural Implant вы разберетесь в прунинге с использованием методов второго порядка.
Но перед этим мы напомним, как устроены «классические» алгоритмы.
Краткий экскурс в историю прунинга
Что такое прунинг?
Вообще, о прунинге нейросетей заговорили, как только появились сами нейросети. Первые заметки есть уже у ЛеКуна. Прунинг — это процедура «выбрасывания» (происходит посредством зануления) ненужных весовых коэффициентов (неструктурированный прунинг) и/или ненужных фильтров/нейронов (структурированный прунинг).
Наглядное изображение структурированного и неструктурированного прунинга
Наглядное изображение структурированного и неструктурированного прунинга
Основное отличие неструктурированного прунинга от структурированного — реальная скорость работы нейронки после всех этих мероприятий. Посудите сами: допустим, у вас была матрица полносвязного слоя NxN. Вы занулили в ней N значений случайным образом. Скорость работы слоя при таком прунинге не поменяется, так как эти нули тоже будут участвовать в перемножениях. Но, если вы занулите целую строку или столбец, то при перемножении его можно пропустить и сеть станет работать чуть быстрее.
Неструктурированный прунинг имеет смысл только тогда, когда у вас есть библиотеки для разреженных вычислений и очень большой уровень прореживания.
Критерии прунинга
Предположим, мы определились, как хотим выбрасывать веса из нейронной сети. Теперь осталось понять, какие веса выбрасывать, то есть рассмотреть различные критерии прунинга.
- Самый простой и понятный критерий — норма фильтра или конкретного веса. Чем меньше вес, тем меньше он вносит вклад в итоговый результат после работы слоя и, тем более вероятно, его можно выбросить. На основе этого критерия проверялась так называемая Lottery Ticket Hypothesis.
- Коэффициент масштабирования γ у BatchNorm. Чем меньше это значение, тем меньше результат на выходе конкретного канала в последующей карте признаков и, тем более вероятно, что его можно выбросить. Идеально, чтобы начать использовать структурированный прунинг и особо не заморачиваться.
- Группа критериев на основе разложения Тейлора. На мой взгляд, это самый интересный подход, и ниже мы разберем его подробнее.
Разложение нейронной сети в ряд
Давайте вспомним, что такое разложение в ряд Тейлора для обычной скалярной функции одной переменной:
Поскольку нейронная сеть — функция многих переменных (как входных данных, так и весов), то давайте возьмем обученную нейронку и зафиксируем входной набор данных. Затем применим к ней функцию потерь и получившуюся конструкцию разложим в ряд. В качестве точки разложения возьмем исходное значение какого-нибудь весового коэффициента, а в качестве желаемого значения возьмем 0.
Потом перенесем всё, что не касается производных в одну сторону и увидим, что получилась разность между значениями функции потерь с обычным и зануленным весом w (L(w=0) — L(w=w_0)). Так как мы не хотим, чтобы функция менялась, берём модуль и получаем удобный критерий прунинга.
\( L(W) = \sum_{p=0}^{P} \frac{L^{(p)}(w_0)}{p!} (w — w_0)^p + R_p(x) \)
\( L(w = 0) = L(w_0) — w_0 \frac{\partial L} {\partial w}{(w_0)} + R_p(W) \)
\( \min |L(w = 0) — L(w)_0 |= \min |w_0 \frac{\partial L}{\partial w}(w_0)| \)
Нижняя строка — так называемый saliency.
Удобен он тем, что у вас есть абсолютно всё для вычисления этого критерия.
- Берете обученную сеть;
- Прогоняете батч на forward и backward;
- Повторяете, пока не накопите некоторое количество градиентов;
- Умножаете вес на градиент, который ему соответствует, сортируете и выбрасываете наименьший по критерию.
Пункт 3 важен, поскольку если вы будете копить градиенты (калиброваться) не на тех данных, то на проде вас ждут крайне неприятные последствия.
💡 Пример: у вас есть сеть, которая распознает кошек и актера Рона Перлмана, исполнившего главную роль в кинофильме Хеллбой. Если вы никогда не подавали семплы с котами на каллибрацию, то вполне логично, что часть связей, ответственных за их распознавание, запрунится. И когда встретится кот, похожий на Рона Перлмана на проде — не удивляйтесь, что ваша нейронка распознает его неправильно.
Разбор статьи
Ну, а теперь, когда у нас есть предварительные знания о прунинге, предлагаю перейти к обзору самой статьи. Из названия понятно, что авторы заложили две идеи: Hessian-aware pruning and Optimal Neural Implant — давайте разберемся с первой.
Hessian-aware pruning
Основная идея в следующем: поскольку нейросеть уже обучена, то она находится в некотором локальном минимуме, а значит, первые производные у нее по идее должны равняться нулю, значит, необходимо учитывать вторые порядки при разложении:
\( \Delta L = L(w + \Delta w) — L(w) = g^T \Delta w + \frac{1}{2} \Delta w^T H \Delta w + O(\| \Delta w \| ^3) \)
Выбрасывая первые порядки, требуя минимизации ΔL и используя метод Лагранжа, получаем:
\( \min_{\Delta w} \frac{1}{2}\Delta w^T H \Delta w = \frac{1}{2} \begin{pmatrix} \Delta w_p \\ \Delta w_l \end{pmatrix}^T \begin{pmatrix} H_{p,p} & H_{p,l} \\ H_{l,p} & H_{l,l} \end{pmatrix} \begin{pmatrix}\Delta w_p \\ \Delta w_l\end{pmatrix}, \\ \text s. \text t. \; \Delta w_p + w_p = 0. \)
Нижние подписи p, l означают запруненные и оставленные параметры соответственно. Применили метод Лагранжа:
\( \frac{1}{2}\Delta w^T H \Delta w = \frac{1}{2} w_p^T (H_{p,p} — H_{p,l} H_{l,l}^{-1} H_{l,p})w_p \)
Если честно, все это высчитывать у нас не хватит никаких вычислительных ресурсов, поэтому в таких случаях строят различные предположения-пожелания:
- Хотим структурированный прунинг, а значит будем выбрасывать группы параметров, которые относятся к одному фильтру/нейрону в слое;
- Слои в нейронной сети у нас независимы. Это конечно не совсем так, особенно, когда мы добавили skip-connection, но допустим.
Из 1 и 2 вытекает то, что гессиан у нас блочно-диагональный. А далее авторы делают хитрый ход. Они говорят: «А давайте каждый из этих блоков аппроксимируем следом!».
При этом, используя такую аппроксимацию, авторы неявным образом делают предположение о том, что веса внутри группы независимы. Ведь след обычного гессиана, блочно-диагонального и строго диагонального — одинаков. Собирая все вместе, получаем:
\( \frac{1}{2}\Delta w^TH\Delta w = \frac{1}{2}w_p^T[H^{-1}]_{p,p}^{-1}w_p \approx \frac{1}{2}w_p^T\frac{Trace(H_{p,p})}{ p}w_p = \frac{Trace(H_{p,p})}{2p}{||w_p||_{2}^2} \)
След гессиана очень удобно считать через так называемый Hutchinson trace estimator:
\( Tr(A) = \mathbb{E}[v^T Av] = \frac{1}{V} \sum_{i=1}^{V} v_i^T Av_i \)
A — матрица, след которой мы хотим посчитать, v_{i} — это некоторый случайный вектор из распределения Радемахера (+1/-1 с вероятностью 0.5).
На мой взгляд все эти трюки с Хатчинсоном, блочной диагональностью и т.д. сделаны только с одной целью: в PyTorch есть только одна нормальная функция, при помощи которой можно посчитать гессиан 😅 Авторы статьи, в свою очередь, прознали про нее, поняли, что вторые порядки проще посчитать именно так и решили подогнать под этот способ вычисления то, что нам надо 🙃
Neural Optimal Implant
Блеск и нищета всех методов прунинга заключается в следующем: мало кто измеряет реальное время работы сети!
Обычно все измеряют так называемые MAC-и (количество операций сложений-умножений), которое сами называют FLOPs-ами. При этом MAC зачастую не отображают реальную картину того, как быстро работает сеть.
💡 Пример: допустим, у нас есть сеть с неструктурированным прунингом. Мы взяли FC-слой NxN и запрунили его так, что осталось всего N значений по диагонали.
Тогда, при подсчете MAC-ов, можно будет сказать, что у нас будет всего 2N-операций сложений-умножений, если на вход придет неразреженная матрица.
Но при этом мы совсем не учитываем, как будут проводиться вычисления в PyTorch или CUDA/CuDNN на самом деле — возможно, и вся матрица перемножается, не смотря на нули.
Вот и получается, что МАС-ов формально 2N, а слой работает долго.
Ну, это была присказка, а теперь сказка.
Собственно, вот мы взяли, применили метод из предыдущего пункта, получили какие-то значения критериев. Идея авторов состоит в том, что не стоит делить каналы только на те, которые вы оставляете и выбрасываете. Среди оставшихся, можно выделить те, которые чуть хуже, чем остальные. Их предлагается «слегка» модифицировать, чтобы поджать их «скорость работы», но при этом совсем не выбрасывать, иначе упадет точность.
Авторы предлагают такие каналы не выбрасывать, а заменить на что-то c меньшей вычислительной сложностью. В своей версии авторы предлагают заменить такие каналы на свертку 1х1. Процедуру можно описать так:
- Выбираем конкретный слой;
- Выбираем в весовых фильтрах этого слоя полезные каналы, «ни туда, ни сюда» и бесполезные;
- Полезные каналы оставляем;
- Бесполезные каналы выбрасываем;
- Те которые ни туда ни сюда — заменяем на свертку 1х1;
- По итогу слой представляет из себя сконкатенированые результаты работы сверток хороших каналов и 1х1.
При этом, насколько это будет эффективно с точки зрения реального времени работы на GPU — непонятно. Да и вычислительный граф будет страшным, как смертный грех, что заставляет беспокоиться об эффективности этого метода.
Вопрос для самых смекалистых подписчиков: как вы думаете, что используют авторы для замеров времени? Напишите в комментариях, за первые 2 верных ответа подарим пиццу!🍕
Результаты
Как всегда, авторы всех побили и победили с огромнейшим отрывом! 🙌
Секрет китайской шкатулки. Смотрим код
Ну, а теперь насладимся великолепным кодом от разработчиков и поймем, почему же все-таки был выбран настолько хитрый способ вычислений гессианов.
Вот код Hutchinson Trace Estimator с репозитория авторов:
v = [torch.randint_like(p, high=2, device=device).float() * 2 - 1 for p in params]
if use_loader:
print('================I USE LOADER================')
THv = [torch.zeros(p.size()).to(device) for p in params]
for inputs, targets in data:
inputs, targets = inputs.to(device), targets.to(device)
model.zero_grad()
outputs = model (inputs)
loss = criterion(outputs, targets)
loss.backward(create_graph=True)
params, gradsH = get_params_grad(model)
Hv = torch.autograd.grad(gradsH, params, grad_outputs=v, only_inputs=True, retain_graph=False)
Обратите внимание на последнюю строку. Тут-то и происходит вся магия. Функция torch.autograd.grad принимает на вход три типа тензоров:
- gradsH — это градиенты, которые мы посчитали в первом прогоне;
- params — тензоры параметров;
- v — те самые случайные векторы из распределения Радемахера.
При обычном использовании, вместо grad_outputs должны стоять градиенты с предыдущего слоя сети, а сама функция torch.autograd.grad используется для back-prop алгоритма. При этом следует обратить внимание, что torch.autograd.grad считает производные как бы поочередно и число их равно числу тензоров из grad_outputs:
\( \frac{\partial gradsH{[i]}}{\partial param{[i]}} * v{[i]} \)
Тут-то и проявляется «независимость» между слоями, которую мы постулировали раньше (мы просто не можем по-другому посчитать 🙃).
После того, как мы посчитали \( H_v \), необходимо умножить еще на \( v \) слева, что они и делают, попутно разбивая на каналы:
with torch.no_grad():
if channelwise:
for Hv_i in range(len(Hv)):
for channel_i in range(Hv[HV_i].size(0)):
trace_vhv[Hv_i][channel_i].append(Hv[Hv_i][channel_i].flatten().dot(v[Hv_i][channel_i].flatten()).item())
Есть альтернативный способ посчитать гессиан, это функция torch.autograd.functional.hessian. Пользоваться ей, конечно же, невозможно, ибо придется выполнять упражнения, связанные с объединением сети и loss-а в одну функцию, где веса — ее входные значения. Именно поэтому авторы статьи пользуются torch.autograd.grad, даже с учетом ее ограничений.
Минутка юмора:
- По умолчанию, у авторов оценка матожидания идет по 3 семплам, градиент считается по одному батчу;
- Авторы соблюдают максимальный баланс.
def update_indices(model, network):
print("updating indices")
dependencies = get_layer_dependencies(model, network)
update_out_indices(model, dependencies)
update_in_dinces(dependencies)
Что делает функция update_out_indices? Разумеется, ничего. Зачем добавили? Видимо, если есть update_in_indices, то должна быть и update_out_indices, ведь инь не может существовать без янь.
def update_indices(model, network) :
print ("updating indices")
dependencies = get_layer_dependencies(model, network)
update_out_indices(model, dependencies)
update_in_dinces(dependencies)
def update_out_indices(model, dependencies):
pass
Вместо заключения
В целом, методы второго порядка это крайне перспективны! Однако очень ограничены в использовании.
- Они слабо применимы при итеративном прунинге (когда прунят веса по чуть-чуть, а потом доучивают), который дает лучшую точность по сравнению с one-shot процедурой (все и сразу до желаемых MACs);
- У методов второго порядка при корректном использовании (а не на одном батче с тремя семплами v) время работы должно быть значительно дольше, чем у тех же методов первого порядка, которые, к тому же можно встроить в процедуру обучения, накапливая градиенты каким-нибудь скользящим средним.
Но зато они учитывают кривизну функции потерь нейронной сети, что в перспективе должно быть лучше с точки зрения генерализации сети!
