Может ли pytorch оптимизировать последовательные операции (например, граф тензорного потока или JAX jit)?

Question

Изначально tenorflow и pytorch имели принципиальное отличие:

• tenorflow основан на вычислительном графе. Построение этого графика и его оценка в сеансе - два отдельных шага. В процессе использования график не меняется, что позволяет проводить оптимизацию.
• torch охотно оценивает операции над тензором. Это делает API более удобным (без сеансов), но также теряет возможность распознавать и оптимизировать операции, которые всегда выполняются последовательно.

Теперь эта разница становится менее очевидной. Tensorflow ответил популярность факела с ТФОМ нетерпеливым. Существует также проект JAX, который основан на той же базовой структуре, что и tensorflow (XLA). JAX не имеет понятия о сеансе. Но он позволяет скомпилировать несколько операций вместе, просто вызвав jit.

Поскольку Tensorflow перешел на функциональность PyTorch, работает ли PyTorch над интеграцией преимуществ Tensorflow? Есть ли в PyTorch (или в его дорожной карте) что-то вроде сессий или jit-функций?

В документации по API есть раздел jit, но, насколько я понимаю, это больше касается экспорта ваших моделей.

Answer 1

Как вы упомянули, есть torch.jit и его целью также является введение оптимизации в экспортируемый граф (например, объединение ядра, оптимизация констант и т. д.). IIRC вы можете найти исходный код в отношении тех, кто в их GitHub репо здесь, хотя я не уверен в том, те, которые явно упоминается где - то в документации (или достаточно явно следует помнить).

поскольку 1.3также введено квантование (см. здесь некоторое введение). В разделе руководств, а именно здесь вы можете увидеть явное слияниеConv2d, BatchNorm а также ReLUв целях повышения производительности. Ofc также существуют определенные вещи, такие как использованиеint вместо того float для весов (квантование), смешанная арифметика (с использованием halfточность с плавающей запятой, когда это возможно, см. NVidia Apex) и другие.

И последнее, но не менее важное: я не думаю, что хорошо написанная модель использует векторизованные операции и экспортируется с torchscriptвы увидите действительно существенные различия во время выполнения из-за некоторой общей оптимизации графа. Тем не менее, все зависит от того, собираетесь ли вы использовать GPU, CPU, TPU, каковы их версии, хотите ли вы использовать только логический вывод или обучение и т. Д. Довольно сложно определить, насколько быстроtensorflow по сравнению с pytorch(помимо некоторых хорошо известных проблем в обеих средах). В общем, это зависит, и измерения сильно различаются AFAIK.

Кстати. Что касается преимуществ каждого фреймворка, их ядро ​​действительно начинает охватывать похожие вещи (в последнее время PyTorch получил поддержку мобильных устройств, см. Здесь). Настоящая разница по-прежнему заключается в другом базовом подходе и в том, что каждая структура должна делать, чтобы обойти эти ограничения.