torch.func

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

• «Преобразование функции» — это функция высшего порядка, которая принимает числовую функцию и возвращает новую функцию, вычисляющую другую величину.
• 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
  • Утилиты отладки
• Ограничения UX
  • Общие ограничения
  • API torch.autograd
  • Ограничения vmap
  • Случайность
• Переход с functorch на torch.func
  • Преобразования функций
  • Утилиты для модулей NN
  • functorch.compile