DeepMind только что выпустила Haiku и RLax для нейронных сетей и обучения с подкреплением

Используйте новейшие библиотеки JAX для облегчения ваших проектов AI

Исследования в области искусственного интеллекта (ИИ) развиваются каждый день, как и его инструменты - передовая компания DeepMind, занимающаяся искусственным интеллектом, только что выпустила две библиотеки на основе JAX для нейронных сетей (NN) и обучения с подкреплением (RL). , а именно Haiku и RLax, соответственно. Поскольку оба имеют открытый исходный код, мы можем использовать эти новейшие библиотеки для нашей собственной работы с ИИ.

Что такое JAX?

Выпущенный Google в 2018 году, JAX представляет собой инструмент с открытым исходным кодом, который генерирует высокопроизводительный код ускорителя путем преобразования программ Python и NumPy. За кулисами JAX хорошо интегрируется Autograd, который может эффективно различать собственные функции Python и NumPy, и XLA, компилятор линейной алгебры, который может оптимизировать модели TensorFlow.

Таким образом, JAX был создан для высокопроизводительных исследований в области машинного обучения, предоставляя интегрированную систему для оптимизации модели обучения. Его популярность растет среди исследователей искусственного интеллекта. Однако это не самая простая вещь в работе, что, вероятно, привело DeepMind к разработке Haiku и RLax.

Что такое Haiku и RLax?

Haiku - это простая библиотека NN для JAX, и эта библиотека разработана некоторыми из авторов Sonnet - библиотеки нейронной сети для TensorFlow. Эта библиотека позволяет исследователям ИИ получить полный доступ к преобразованиям чистых функций JAX, предоставляя модели объектно-ориентированного программирования. Предоставляя два основных инструмента - hk.Module и hk.transform, Haiku предназначен для управления параметрами модели NN и другими состояниями модели.

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

Почему они?

Поскольку обе библиотеки стали доступны с открытым исходным кодом, у разработчиков может возникнуть соблазн рассмотреть возможность их использования в своих собственных проектах. Конечно, вероятно, еще преждевременно использовать их в производственном проекте, но есть два основных фактора, которые могут оправдать то, что мы должны начать с ними экспериментировать.

  • Общее высокое качество. Обе библиотеки поддерживаются DeepMind и были протестированы в разумных масштабах. Например, авторы Haiku также разрабатывают Sonnet, уважаемую библиотеку TensorFlow для NN.
  • Удобство для пользователя. Обе библиотеки выполняют набор четко определенных простых в использовании операций. Для Haiku его API и абстракции очень близки к Sonnet, что делает переход на Haiku безболезненным. Если вы не знакомы с Sonnet, вы все равно можете научиться использовать его относительно быстро из-за наличия в нем объектно-ориентированных моделей программирования.

Выводы

Обе библиотеки облегчат разработку моделей NN и ML, предоставляя отличный инструмент интеграции в рамках JAX. Если вы хотите попробовать что-то новое, я определенно рекомендую вам рассмотреть возможность использования этих библиотек, по крайней мере, для экспериментальных целей.

Полезные ссылки

GitHub JAX

GitHub Haiku

GitHub RLax