Что такое граф вычислений (computational graph) в PyTorch и как он работает?

Граф вычислений в PyTorch — это динамическая (define-by-run) направленная ациклическая графовая структура (DAG), которая автоматически строится в памяти во время выполнения операций над тензорами с включенным отслеживанием градиентов (requires_grad=True). Его основная цель — обеспечить автоматическое дифференцирование (autograd) для вычисления градиентов методом обратного распространения ошибки (backpropagation).

Ключевые компоненты графа:

• Узлы (Nodes): Представляют собой тензоры. Каждый узел хранит:
  • Данные (.data).
  • Градиент (.grad), если он был вычислен.
  • Ссылку на функцию (.grad_fn), которая его создала (запись операции).
• Ребра (Edges): Определяют поток данных от входных тензоров к выходным через операции.

Как это работает на практике:

import torch

# 1. Создание тензоров с отслеживанием градиентов
x = torch.tensor(2.0, requires_grad=True)
w = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)

# 2. Выполнение операций. На этом этапе ПОСТРОЕНИЯ ГРАФА.
# Каждая операция записывается в граф.
y = w * x + b  # y = 3*2 + 1 = 7
# Теперь y.grad_fn указывает на объект AddBackward,
# который, в свою очередь, ссылается на MulBackward для w*x.

# 3. Инициирование обратного распространения
loss = (y - 5) ** 2  # Допустим, наша цель была 5, loss = (7-5)^2 = 4
loss.backward()  # Запуск backpropagation по графу

# 4. Градиенты вычислены и сохранены в leaf-тензорах
print(f"Градиент d(loss)/dx: {x.grad}")  # 2*(y-5)*w = 2*2*3 = 12
print(f"Градиент d(loss)/dw: {w.grad}")  # 2*(y-5)*x = 2*2*2 = 8
print(f"Градиент d(loss)/db: {b.grad}")  # 2*(y-5)*1 = 4

Важные особенности PyTorch:

1. Динамичность (Eager Execution): Граф строится «на лету» по мере выполнения кода. Это позволяет использовать стандартные конструкции Python (циклы, условия) внутри модели.
2. Эффективность памяти: По умолчанию граф автоматически удаляется после вызова .backward(), чтобы освободить память. Для многократного использования графа можно передать retain_graph=True.
3. Контроль отслеживания: Контекстные менеджеры отключают построение графа для инференса и экономии памяти:

   with torch.no_grad():
       inference_output = model(x)  # Граф не строится, быстрее и меньше памяти

4. JIT-компиляция (TorchScript): Динамический граф можно «заморозить» и скомпилировать в статический для продакшн-деплоя.

Итог: Граф вычислений — это фундаментальный механизм PyTorch, который делает возможным удобное и гибкое обучение нейронных сетей с автоматическим вычислением градиентов.