Привет! На связи команда «БАРС Груп», и мы продолжаем разговор о фреймворке PyTorch.

PyTorch — это фреймворк ML для Python с открытым исходным кодом, широко применяемый для решения прикладных задач, связанных с нейросетями. Как правило, фреймворки машинного обучения часто заточены либо на удобство использования, либо на скорость. PyTorch же отличается тем, что сочетает в себе оба преимущества. Он поддерживает код как модель, упрощает отладку и согласуется с другими популярными библиотеками научных вычислений, оставаясь при этом эффективным и поддерживая аппаратные ускорители, такие как графические процессоры. При этом каждый аспект PyTorch — это обычная программа Python, находящаяся под полным контролем пользователя.

Это вторая часть статьи‑перевода от команды разработчиков PyTorch (Адама Пашке, Сэма Гросса и их единомышленников). В первой части авторы разобрали принципиальные отличия PyTorch от DyNet и других фреймворков и библиотек автоматического дифференцирования, а также особенности его интерфейса (флаги переменных, хуки, расширения). Сегодня — информация о реализации данного фреймворка в таких аспектах как управление памятью (оперативная очистка промежуточных значений, когда они становятся ненужными), выполнение операций над тензором и способ их аннулирования.

Реализация

Внутри Переменная (Variable) представляет собой просто обертку вокруг тензора (Tensor), которая также содержит ссылку на граф объектов Функций (Functions). Этот граф является неизменным, чисто функциональным представлением производной вычисляемой функции. Переменные — это просто изменяемые указатели на этот граф (они изменяются, когда происходит операция на месте).

Функции можно рассматривать как замыкания, содержащие весь контекст, необходимый для вычисления векторно‑якобианских произведений. Они принимают градиенты выходных данных и возвращают градиенты входных данных (формально левое произведение left product, включающее выражение для соответствующей операции). Граф функций — это замыкание с одним аргументом, которое принимает левое произведение (left‑product) и умножает его на производные от всех содержащихся в нем операций. Передаваемые левые произведения сами по себе являются переменными, что делает оценку графа дифференцируемой.

Управление памятью

Основной вариант использования PyTorch — обучение моделей машинного обучения на графическом процессоре. Поскольку одно из самых больших ограничений графических процессоров — это небольшой объем памяти, PyTorch уделяет большое внимание тому, чтобы все промежуточные значения освобождались, как только они становятся ненужными. Действительно, Python хорошо подходит для этой цели, потому что он по умолчанию подсчитывает ссылки (используя сборщик мусора только для прерывания циклов).

Переменная (Variable) и Функция (Function) в PyTorch должны быть спроектированы так, чтобы хорошо работать в режиме подсчета ссылок. Например, Функция записывает указатели на Функцию, которая принимает ее результаты, так что подграф функций освобождается, когда его сохраняющая выходная Переменная становится неиспользуемой. Это противоположно обычному владению замыканиями, когда замыкание сохраняет вызываемые им другие замыкания.

Другая проблема заключается в избегании циклов в графе ссылок. Простая реализация автоматического дифференцирования может легко вводить такие циклы (например, когда дифференцируемая функция хочет сохранить ссылку на свой вывод). PyTorch разбивает их, записывая не полноценную переменную, а «сохраненную переменную», которая в таких случаях не содержит указатель непосредственно на саму функцию.

Операторы C++

Несмотря на то, что все операции в Python можно выразить с помощью API расширений, они требуют больших накладных расходов на интерпретатор. Перенос операторов на C++ снижает накладные расходы и снижает задержку одной дифференцируемой операции, отправляемой из Python, до 3,4 мкс по сравнению с 1,7 мкс для тензорной операции. Дополнительным преимуществом является то, что можно иметь несколько потоков, выполняющих их параллельно (в отличие от Python, который ограничивает параллелизм из‑за GIL). Это особенно важно в контексте нескольких графических процессоров, которые не могут быть загружены одним потоком центрального процессора.

Поддержка операций на месте

Часто пользователи PyTorch хотят выполнять операции с тензором на месте, чтобы избежать выделения нового тензора, когда известно, что он не нужен. Интуитивно понятно, что операция на месте эквивалентна соответствующей операции не на месте. Исключение — когда переменная, которая изменяется на месте, имеет свою историю вычислений, «перебазированную» так, чтобы она указывала на производную оператора на месте, вместо указания на предыдущую функцию (история вычислений всегда остается чисто функциональной). Однако эти операции на месте тонко взаимодействуют с autograd.

Интуитивно понятно, что операция на месте эквивалентна соответствующей операции не на месте, за исключением того, что переменная, которая с на месте, имеет «перебазированную» историю вычислений, указывающую на производную оператора на месте, а не на его предыдущую функцию (история вычислений всегда остается чисто функциональной). Однако эти операции на месте тонким образом взаимодействуют с механизмами автодифференцирования.

Аннулирование

Операция на месте может сделать недействительными данные, необходимые для вычисления производных. Рассмотрим следующий пример:

y = x.tanh()

y.add_(3)

y.backward()

Эта программа указывает PyTorch выполнить операцию на месте с y, однако это неверно, если y, значение которого равно tanh(x), было сохранено, чтобы его можно было использовать в обратном вычислении (напомним, что tanh» (x) = 1 − tanh2 (x)).

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

Псевдонимы

PyTorch поддерживает нетривиальные псевдонимы между переменными. Такие операции, как транспонирование и сужение, создают новые тензоры с новыми размерами и шагами, которые имеют общий доступ к хранилищу с исходными тензорами. Проблема с псевдонимами заключается в том, что они могут потребовать нетривиальных преобразований истории вычислений многих переменных. Рассмотрим следующий пример:

y = x[:2]

x.add_(3)

y.backward()

Обычно операция на месте над x влияет только на историю вычислений x. Однако в этом случае добавление на месте к x также приводит к обновлению некоторых элементов y. Таким образом, история вычислений y также изменилась. Поддержка этого случая довольно нетривиальна, поэтому PyTorch отклоняет эту программу, используя дополнительное поле в счетчике версий (см. Аннулирование выше), чтобы определить, что данные являются общими.

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

В заключение — несколько ссылок для тех, кто хочет поближе узнать PyTorch и изучить тему автоматического дифференцирования на данном фреймворке:

книга «Глубокое обучение на Python» (переведена на русский язык, издательство МИФ), выжимка книги доступна на Хабре);
статья «PyTorch — ваш новый фреймворк глубокого обучения»;
статья Адама Пашке «Pytorch: императивный стиль, высокопроизводительная библиотека глубокого обучения» (на английском языке).

Особенности автоматического дифференцирования в PyTorch. Часть 2