torch.func

torch.func, ранее известный как «functorch», представляет собой аналогичные JAX композируемые преобразования функций для PyTorch.

Что такое композируемые преобразования функций?

• «Преобразование функции» — это функция высшего порядка, которая принимает численную функцию и возвращает новую функцию, вычисляющую другую величину.
• torch.func имеет преобразования автоматического дифференцирования (grad(f) возвращает функцию, вычисляющую градиент f), преобразование векторизации/пакетирования (vmap(f) возвращает функцию, вычисляющую f по пакетам входных данных) и другие.
• Эти преобразования функций могут комбинироваться друг с другом произвольно. Например, комбинирование vmap(grad(f)) вычисляет величину, называемую градиентами по каждому образцу, которую PyTorch сегодня не может эффективно вычислить.

Зачем нужны композируемые преобразования функций?

Существует ряд сценариев использования, которые трудно реализовать в PyTorch сегодня:

• вычисление градиентов по каждому образцу (или других величин по каждому образцу)
• запуск ансамблей моделей на одном компьютере
• эффективное объединение задач в цикле внутреннего уровня MAML
• эффективное вычисление Якобианов и Гессианов
• эффективное вычисление пакетированных Якобианов и Гессианов

Комбинирование vmap(), grad() и vjp() преобразований позволяет нам выразить вышеперечисленное, не разрабатывая отдельную подсистему для каждого случая. Эта идея композируемых преобразований функций заимствована из фреймворка JAX.

Подробнее

• Быстрый обзор torch.func
  • Что такое torch.func?
  • Зачем нужны композируемые преобразования функций?
  • Какие преобразования существуют?
• Справочник API torch.func
  • Преобразования функций
  • Утилиты для работы с torch.nn.Modules
• Ограничения в пользовательском интерфейсе
  • Общие ограничения
  • API torch.autograd
  • Ограничения vmap
  • Случайность
• Миграция с functorch на torch.func
  • Преобразования функций
  • Утилиты для модулей NN
  • functorch.compile