Работа с контейнерами Nx

`Nx.Container` для вложенных структур данных в Elixir

В численных вычислениях распространенной задачей является необходимость работы с наборами тензоров. Эти коллекции могут иметь форму кортежа, карты или другой, возможно, вложенной структуры данных. Например, вы можете захотеть сохранить параметры такой модели, как нейронная сеть, на карте, чтобы было легко определить, какие параметры относятся к каким слоям. Nx упрощает процесс работы с произвольно вложенными контейнерами с помощью Nx.Container API.

Если вы знакомы с экосистемой машинного обучения Python, возможно, вы знакомы с tf.nest API или pytrees JAX. И tf.nest, и pytrees создают абстракции для работы с древовидными структурами данных. В обоих API тензоры обрабатываются как значения листа или атома. Такие функции, как tf.nest.map_structure и tree_map в JAX, работают аналогично Enum.map в Elixir, но они применяют функции к атомарным значениям во вложенной структуре, сохраняя исходную древовидную структуру. Например, если у вас есть коллекция:

{'foo': (1, {'bar': 2, 'baz': (3, 4, (5,))})}

Затем применение map_structure или tree_map к вложенной структуре с функцией x * x приведет к новой структуре:

{'foo': (1, {'bar': 4, 'baz': (9, 16, (25,))})}

Nx предлагает аналогичный API с Nx.Container . Nx.Container позволяет пользователям перемещаться по произвольно вложенным контейнерам и регистрировать пользовательские структуры данных в качестве контейнеров для использования с числовыми определениями Nx. Каждый контейнер Nx должен реализовать две функции или @derive реализацию стандартного Nx.Container протокола. В протоколе есть две функции reduce и traverse, которые имеют следующие подписи:

reduce(t(), acc, (Nx.Tensor.t(), acc -> acc)) :: acc
traverse(t(), acc, (Nx.Tensor.t(), acc -> {term(), acc})) :: acc

reduce применяет функцию к конечным значениям внутри контейнера, перенося и возвращая аккумулятор. traverse делает то же самое; однако он также возвращает обновленные конечные значения вместе с обновленным аккумулятором. Важно отметить, что обе эти функции не применяют функции вложенным образом, а возлагают ответственность за обработку вложенных обходов на пользователя. Учитывая эти две функции, вы можете реализовать аналогичные функции для tree_map или tf.nest.map_structure:

defn tree_map(container, fun) do
  {container, :ok} = Nx.Container.traverse(container, :ok, &do_tree_map(&1, &2, fun)
end
defp do_tree_map(container_or_tensor, :ok, fun) do
  case container_or_tensor do
    %Nx.Tensor{} = leaf ->
      {fun.(leaf), :ok}
    container ->
      {tree_map(container, fun), :ok}
  end
end

Вы также можете реализовать более сложные функциональные преобразования, такие как reduce, map_reduce, zip и так далее.

Nx поддерживает реализации контейнеров для карт и кортежей из коробки. Добавить поддержку пользовательского контейнера, такого как пользовательская структура, так же просто, как использовать @derive Elixir для вашей пользовательской структуры:

defmodule MyData do
  
  @derive {Nx.Container,
           containers: [:value_1, :value_2],
           keep: [:option]}
  defstruct [:option, :value_1, :value_2]
end

:keep указывает, какие поля структуры следует сохранить или сохранить, когда структура используется в defn . Это полезно для сохранения метаданных, таких как атомы. :containers укажите, какие поля содержат тензоры (или другие контейнеры). Внедрив Nx.Container для своей структуры, вы можете без проблем использовать пользовательские структуры данных в defn:

defmodule Functions do
  import Nx.Defn
defn function_on_my_data(my_data) do
    operation = get_operation(my_data)
    operation.(my_data.value_1, my_data.value_2)
  end
defnp get_operation(my_data) do
    transform(my_data, fn %MyData{option: op} -> &apply(Nx, op, [&1, &2]) end)
  end
end
data = struct(MyData, %{option: :add, value_1: 1, value_2: 2})
Functions.function_on_my_data(data) |> IO.inspect

Nx.Container — это гибкая абстракция, упрощающая работу с вложенными структурами данных и регистрацию пользовательских структур данных в defn Nx. Для получения дополнительной информации я рекомендую прочитать Документацию Nx и проверить некоторые реализации контейнеров в дикой природе, такие как Axon.

Шон Мориарти и Хосе Валим вместе работали над созданием Библиотеки Nx Эликсира. Шон также написал книгу Генетические алгоритмы в эликсире вместе с The Pragmatic Bookshelf:

Генетические алгоритмы в Эликсире
От финансов до искусственного интеллекта, генетические алгоритмы — это мощный инструмент с широким спектром приложений. Но…pragprog.com

Вы можете сэкономить 35% на электронной версии с промокодом smgaelixir_35 до 31 июля 2022 г.

Другие статьи Шона Мориарти

Короткая викторина Nx
Отточите свои навыки программирования массивов с помощью библиотеки Nx от Elixirmedium.com

Написание генетического алгоритма на Nx
Численные вычисления для Elixirmedium.com

Работа с контейнерами Nx