Потерянная тетрадь — разработка длиною в ночь

Был поздний вечер. Ужасно болела голова, простуда не давала нормально дышать, а градусник упрямо полз вверх. Что можно делать в таком состоянии, когда кажется, что сама Смерть уже натирает лыжи где‑то рядом?Казалось бы, ничего значимого.

Поэтому я просто решил разобрать груду скопившихся вокруг бумаг — надо же иногда наводить порядок. И тут же обратил внимание на старую тетрадь торчавшую из‑за шкафа.Я аккуратно развернул ее, и взгляд зацепился за старый рисунок, который был сделан немногим позже ухода NVIDIA из России, где я работал ведущим инженером и в основном занимался написанием фреймворка для нейросетей, но уезжать не стал. А рисунок набросал, размышляя о том, как бы заполучить заветные ядра свертки Винограда без китайских теорем, универсально, для абсолютно любой формы ядра, включая те, которые не решаются аналитически. Идея была простой: почему бы матрицы Винограда не получить методом градиентного спуска?Я быстро озадачил свой локальный ИИ вопросом: а не сделал ли кто‑то за эти годы то же самое? Успокаивающее шуршание дросселей на старых V100, говорило о том что процесс пошел, а помаргивание коммутатора подтверждало что нейронки активно занимаются поиском ответа… Выскочила пара ссылок на AlphaTensor от DeepMind и ещё несколько академических работ, довольно слабо относящихся к реальной «серебряной пуле». Но главное — ни у кого из них я не увидел двух ключевых вещей: открытого кода и универсальности подхода для построения точных алгоритмов под ЛЮБЫЕ свертки (с произвольным Stride, Dilation и асимметрией), но с минимальным числом умножений, как это делал тот крутой мужик, которого звали Шмуэль Виноград. Гештальт, висевший в воздухе годами, потребовал срочной развязки…. Скажу честно: я пишу в основном на C++, CUDA и HDL (VHDL/SystemVerilog). Всё остальное без серьезной надобности не употребляю. Но сегодня спуститься до уровня Python, думаю, было вполне простительно — держать стандартную дисциплину разработки я был попросту не в состоянии.Поэтому я просто попросил ИИ набросать мне каркас нейросети прямо по моему рисунку из тетради. Модель справилась легко — архитектура‑то концептуально очень простая.Два полносвязных слоя (fc_D и fc_F) кодируют входные данные D и ядро F в промежуточное пространственное «горлышко» (inter_dim). В горлышке происходит поэлементное умножение​. Третий полносвязный слой (fc_C) декодирует результат обратно в форму выхода. Дальше я выстроил простейший цикл обучения, где источником «истины» и входных данных служила сама эталонная свертка (F.conv1d, F.conv2d или F.conv3d), которую мы хотим сжать. Генерируем случайные нормальные батчи D и F, считаем от них истинную свертку — вот вам и идеальный бесконечный датасет с нулевым шумом. Чтобы заставить сеть обнулять лишние связи в матрицах (ведь мы ищем разреженность, как у Винограда), я повесил на веса обычный L1-штраф и небольшой прунинг по теоретически минимальной дроби: Натравил Adam, настроил простейшую 1D‑свертку, нажал Run и совершенно спокойно приготовился ждать провала операции…И когда на экране выкристаллизовалась совершенно ясная, чистая матрица! Моя реакция была: «Да ладно?!», как когда‑то у моего напарника по лабораторке, когда я написал сотни строк кода на ассемблере, и они заработали строго с первого раза. Итак, вечер переставал быть стремным. А кто‑то там в темноте, кажется, тихо снял лыжи, отложил косу и подсел поближе… Конечно, сразу захотелось попробовать геометрии побольше. Алгоритм очевидно сходился, но делал это довольно неторопливо. Модели…

Источник: habr.com — habr.com/ru/articles/1093580/