Обучение нейронных сетей для Leela Zero с помощью PyTorch и PyTorch Lightning

Простой обучающий конвейер для Leela Zero, реализованный с помощью PyTorch, PyTorch Lightning и Hydra.

Недавно я искал способы ускорить свои исследования и управлять своими экспериментами, особенно в отношении написания обучающих конвейеров и управления конфигурациями экспериментов, и я обнаружил эти два новых проекта под названием PyTorch Lightning и Hydra. PyTorch Lightning помогает быстро писать обучающие конвейеры, а Hydra помогает аккуратно управлять конфигурациями.

Чтобы попрактиковаться в их использовании в более реалистичной обстановке, я решил написать обучающий конвейер для Leela Zero, движка Go. Я выбрал это, потому что это масштабный проект с интересными техническими проблемами, связанными с обучением огромных сетей на больших наборах данных с использованием нескольких графических процессоров. Кроме того, раньше мне нравилось реализовывать уменьшенную версию AlphaGo для шахмат, поэтому я подумал, что это будет забавный побочный проект.

В этом блоге я объясню основные детали проекта, чтобы вы могли легко понять, что я сделал. Вы можете прочитать мой код здесь: https://github.com/yukw777/leela-zero-pytorch

Лила Зеро

Первым шагом было выяснить, как устроена нейронная сеть Лилы Зеро. Я сильно сослался на документацию Лилы Зеро и ее обучающую программу Tensorflow.

Архитектура нейронной сети

Нейронная сеть Лилы Зеро состоит из башни ResNet с двумя головами, заголовком политики и заголовком значения, как описано в статье AlphaGo Zero. Все фильтры свертки имеют размер 3x3, за исключением фильтров в начале политики и заголовка значения, которые имеют размер 1x1, как в документе. Характеристики игры и доски кодируются как тензоры формы [размер партии, ширина доски, высота доски, количество функций] и сначала передаются через башню ResNet. Затем башня извлекает абстрактные признаки и передает их через каждую из голов, чтобы вычислить распределение вероятностей политики для следующего хода и ценность игры, чтобы предсказать победителя игры.

Вы можете найти детали реализации сети во фрагменте кода ниже.

Формат веса

Лила Зеро использует простой текстовый файл для сохранения и загрузки веса сети. Каждая строка в текстовом файле имеет ряд чисел, которые представляют веса каждого слоя сети. Сначала идет остаточная башня, за ней следует заголовок политики, а затем заголовок значения.

Сверточные слои имеют 2 ряда весов:

Веса свертки с формой [вывод, ввод, размер фильтра, размер фильтра]
Смещение канала

Слои Batchnorm имеют 2 весовых ряда:

Батчнорм означает
Вариации Батчнорма

Слои внутреннего продукта (полностью связанные) имеют 2 ряда весов:

Вес слоя с формой [вывод, ввод]
Смещения вывода

Я написал модульные тесты, чтобы убедиться, что мои файлы веса верны. Еще одна простая проверка работоспособности, которую я использовал, заключалась в том, чтобы рассчитать количество слоев и сравнить его с тем, что говорит Лила Зеро после загрузки моих файлов с весами. Уравнение количества слоев:

n_layers = 1 (version number) +
           2 (input convolution) + 
           2 (input batch norm) + 
           n_res (number of residual blocks) * 
           8 (first conv + first batch norm + 
              second conv + second batch norm) + 
           2 (policy head convolution) + 
           2 (policy head batch norm) + 
           2 (policy head linear) + 
           2 (value head convolution) + 
           2 (value head batch norm) + 
           2 (value head first linear) + 
           2 (value head second linear)

Пока это кажется достаточно простым, но есть необычная деталь реализации, о которой вам нужно знать. Лила Зеро фактически использует смещение для сверточного слоя, чтобы представить обучаемые параметры (gamma и beta) следующего слоя пакетных норм. Это было сделано для того, чтобы формат файла весов, в котором есть только одна линия для весов слоев и другая для смещения, не менялся при добавлении слоев с пакетными нормами.

В настоящее время Leela Zero использует только beta термин пакетной нормы и устанавливает gamma равным 1. Тогда как вы на самом деле используете сверточное смещение для получения тех же результатов, что и применение обучаемых параметров в пакетной норме? Давайте сначала посмотрим на уравнение для нормы партии:

y = gamma * (x — mean)/sqrt(var — eps) + beta

Поскольку Лила Зеро устанавливает gamma в 1, уравнение принимает следующий вид:

y = (x — mean)/sqrt(var — eps) + beta

Теперь пусть x_conv будет выходом сверточного слоя без смещения. Затем мы хотим добавить некоторую погрешность к x_conv, чтобы при выполнении его через пакетную норму без beta результат был таким же, как при выполнении x_conv через уравнение пакетной нормы только с beta, упомянутым выше. В форме уравнения:

(x_conv + bias — mean)/sqrt(var — eps) = 
(x_conv — mean)/sqrt(var — eps) + beta 
x_conv + bias — mean = 
x_conv — mean + beta * sqrt(var — eps) 
bias = beta * sqrt(var — eps)

Итак, если мы установим сверточное смещение на beta * sqrt(var — eps) в файле весов, мы получим желаемый результат, и именно это делает LeelaZero.

Тогда как мы на самом деле это реализуем? В Tensorflow вы можете указать слою пакетной нормы игнорировать только термин gamma, вызвав tf.layers.batch_normalization(scale=False) и покончив с этим. К сожалению, в PyTorch вы не можете настроить слои пакетной нормализации так, чтобы они игнорировали только gamma; вы можете игнорировать только gamma и beta, задав для параметра affine значение False: BatchNorm2d(out_channels, affine=False). Итак, я установил пакетную нормализацию, чтобы игнорировать оба, а затем просто добавил тензор после, который представляет beta. Затем я использовал уравнение bias = beta * sqrt(var — eps), чтобы вычислить сверточное смещение для файла весов.

Учебный конвейер

Разобравшись в деталях нейронной сети Лилы Зерос, пришло время заняться конвейером обучения. Как я уже упоминал, я хотел попрактиковаться в использовании двух инструментов - PyTorch Lightning и Hydra - для ускорения написания обучающих конвейеров и четкого управления конфигурациями экспериментов. Давайте подробнее рассмотрим, как я их использовал.

PyTorch Lightning

Написание обучающего конвейера - это, безусловно, моя наименее любимая часть исследования: он включает в себя много повторяющегося шаблонного кода и его трудно отлаживать. Из-за этого PyTorch Lightning стал для меня глотком свежего воздуха. Это легкая библиотека без множества вспомогательных абстракций поверх PyTorch, которая берет на себя большую часть стандартного кода при написании обучающих конвейеров. Это позволяет вам сосредоточиться на более интересных частях ваших обучающих конвейеров, таких как архитектура модели, и сделать ваш исследовательский код более модульным и отлаживаемым. Кроме того, он поддерживает обучение с использованием нескольких графических процессоров и TPU прямо из коробки!

Чтобы использовать PyTorch Lightning для моего конвейера обучения, больше всего мне пришлось написать код, который я назвал NetworkLightningModule, который наследуется от LightningModule, чтобы указать детали моего конвейера обучения и передать его Trainer. Вы можете следить за официальной документацией PyTorch Lightning, чтобы узнать, как написать свой собственный LightningModule.

Гидра

Другая часть исследования, в котором я искал хорошее решение, - это управление экспериментами. Когда вы проводите исследование, вы неизбежно запускаете множество вариантов своего эксперимента для проверки своей гипотезы, и чрезвычайно важно отслеживать их масштабируемым образом. До сих пор я полагался на файлы конфигурации для управления вариантами экспериментов, но использование простых файлов конфигурации быстро становится неуправляемым. Шаблоны - одно из решений этой проблемы. Однако я обнаружил, что шаблоны в конечном итоге также становятся беспорядочными, потому что, когда вы накладываете несколько слоев файлов значений для визуализации файлов конфигурации, становится трудно отслеживать, какое значение было получено из какого файла значений.

Hydra, с другой стороны, представляет собой систему управления конфигурацией на основе композиции. Вместо того чтобы иметь отдельные шаблоны и файлы значений для визуализации окончательной конфигурации, вы объединяете несколько меньших файлов конфигурации, чтобы составить окончательную конфигурацию. Она не такая гибкая, как система управления конфигурацией на основе шаблонов, но я считаю, что системы на основе композиции обеспечивают хороший баланс между гибкостью и удобством обслуживания. Hydra - одна из таких систем, специально созданная для исследовательских скриптов. Его вызов немного затруднен, поскольку требует, чтобы вы использовали его в качестве декоратора для функции основной точки входа вашего скрипта, но я действительно думаю, что этот выбор дизайна упрощает интеграцию с вашими учебными скриптами. Кроме того, он позволяет вручную переопределять конфигурации через командную строку, что очень полезно при запуске различных вариантов вашего эксперимента. Я использовал Hydra для управления различными размерами сетевой архитектуры и конфигураций обучающего конвейера.

Оценка

Чтобы оценить свои обученные сети, я использовал GoMill для проведения турниров по го. Это библиотека для проведения турниров между движками Go Text Protocol (GTP), одной из которых является Leela Zero. Вы можете найти конфигурацию турнира, которую я использовал здесь.

Заключение

Используя PyTorch-Lightning и Hydra, я смог значительно ускорить написание обучающих конвейеров и эффективно управлять конфигурациями экспериментов. Я надеюсь, что этот проект и сообщение в блоге также помогут вам в вашем исследовании. Вы можете ознакомиться с кодом здесь: https://github.com/yukw777/leela-zero-pytorch