Этот урок открывает цикл про micrograd, маленькую библиотеку автоматического дифференцирования, вокруг которой построено почти всё, что мы будем делать дальше в курсе. Сегодня мы разберём только половину механики: как число оборачивается в объект Value и как операции над такими объектами незаметно для нас строят граф вычислений. Вторую половину, обратный проход и правило цепочки, оставим на следующий урок. После этого занятия вы сможете читать прямой проход любой формулы в терминах micrograd и самостоятельно пройтись по графу вычислений, который она породила.

Неофициальный курс AI University по открытому коду (MIT). В уроке приводится код из karpathy/micrograd © Andrej Karpathy, лицензия MIT; комментарии переведены на русский, объяснения написаны нашей командой. Курс не связан с автором кода и не одобрен им.

Зачем вообще нужен граф вычислений

Чтобы обучать нейросеть, нужно знать, как функция потерь меняется при малом изменении каждого параметра, то есть нужны частные производные ∂L/∂w по каждому весу w. Есть три способа их получить.

Численное дифференцирование вычисляет производную по определению: (f(x+h) - f(x)) / h при маленьком h. Это просто реализовать, но результат приближённый, зависит от выбора h, и для каждого параметра нужен отдельный проход через всю функцию. Если параметров миллионы, придётся миллионы раз пересчитывать функцию целиком, это неприемлемо медленно.

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

Автоматическое дифференцирование в режиме reverse mode, которое и реализует micrograd, берёт лучшее из обоих миров. Оно не хранит символьную формулу производной целиком, а вместо этого во время прямого прохода запоминает, из каких операций и каких входов получилось каждое промежуточное значение. После того как вычислена финальная величина (например, функция потерь), мы проходим по этой записи в обратном порядке и на каждом шаге применяем правило цепочки, используя только локальную производную конкретной операции. В итоге мы получаем точные производные по всем параметрам сразу за один обратный проход, а не по одному на параметр, как при численном дифференцировании.

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

Класс Value: что он хранит

Откроем определение класса и посмотрим на его конструктор.

# Фрагмент: micrograd/micrograd/engine.py

class Value:
    """ хранит одно скалярное значение и его градиент """

    def __init__(self, data, _children=(), _op=''):
        self.data = data
        self.grad = 0
        # внутренние переменные, нужные для построения графа автоматического дифференцирования
        self._backward = lambda: None
        self._prev = set(_children)
        self._op = _op # операция, породившая этот узел, для graphviz / отладки и т.п.

Разберём поля по отдельности.

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

grad накопитель градиента. Изначально он равен нулю: до обратного прохода мы ещё не знаем, как это значение влияет на итоговый результат. В следующем уроке мы увидим, как сюда записывается ∂L/∂(это значение).

_prev множество родительских узлов, то есть тех Value, из которых получилось текущее значение. Это и есть рёбра графа, только направленные «от результата к входам». Почему именно множество, а не список или кортеж? Во-первых, порядок родителей для построения графа не важен, важен сам факт связи. Во-вторых, множество автоматически избавляет от дублей: если один и тот же узел участвует в операции дважды (как в выражении x * x, где оба операнда это один и тот же объект x), он попадёт в _prev только один раз, а при обходе графа на обратном проходе мы не захотим обработать его дважды.

_op строка с названием операции, которая породила узел: '+', '*', '**2' и так далее. Она ни на что не влияет в вычислениях, это просто подпись для отладки и для визуализации графа.

_backward функция без аргументов, которая умеет доносить градиент из out до его родителей. По умолчанию это lambda: None, пустышка, которая ничего не делает, нужна она для узлов, у которых нет родителей, то есть для «листьев» графа, созданных напрямую из числа. Как эта функция заполняется для каждой операции, разберём подробно в следующем уроке, сегодня нам важно только, что она существует и прикрепляется к каждому новому узлу.

Как операции строят граф: add и mul

Посмотрим на сложение.

# Фрагмент: micrograd/micrograd/engine.py

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), '+')

        def _backward():
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward

        return out

Первая строка решает важную практическую задачу: что если мы напишем a + 5, где a это Value, а 5 обычное число? Тогда other окажется числом int, а не Value, и у него не будет ни .data, ни ._prev. Строка other = other if isinstance(other, Value) else Value(other) оборачивает голое число в Value на лету, создавая для него новый узел без родителей (листовой узел графа). После этого оба операнда гарантированно являются Value, и дальнейший код может не думать о смешанных типах.

Дальше создаётся out, новый узел Value, data которого это обычная сумма self.data + other.data, а родителями указан кортеж (self, other), превращающийся внутри конструктора в множество _prev. Метка операции '+'. На этом прямой проход для сложения закончен: мы получили число-результат и узел, который помнит, что он возник как сумма двух конкретных узлов.

Функция _backward внутри __add__ нам сегодня не нужна в деталях, достаточно заметить, что она определяется отдельно для каждой операции и прикрепляется к out._backward. Как именно она использует out.grad, чтобы разнести градиент по self и other, мы разберём в следующем уроке.

Умножение устроено совершенно симметрично:

# Фрагмент: micrograd/micrograd/engine.py

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), '*')

        def _backward():
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = _backward

        return out

Та же обёртка числа в Value, то же создание out с данными-произведением, родителями (self, other) и меткой '*'. Разница только в формуле для data и в том, какая функция _backward туда положена, но об этом в следующий раз.

Возведение в степень работает немного иначе, потому что показатель степени это не Value, а обычное число (micrograd поддерживает только возведение в постоянную степень, а не Value в степени Value):

# Фрагмент: micrograd/micrograd/engine.py

    def __pow__(self, other):
        assert isinstance(other, (int, float)), "only supporting int/float powers for now"
        out = Value(self.data**other, (self,), f'**{other}')

        def _backward():
            self.grad += (other * self.data**(other-1)) * out.grad
        out._backward = _backward

        return out

Здесь у out только один родитель, self, ведь показатель степени не является узлом графа, он просто число-параметр операции. Метка операции формируется динамически, например '**2' или '**-1', что удобно при отладке.

Функция relu следует тому же шаблону, но содержит ветвление в прямом проходе:

# Фрагмент: micrograd/micrograd/engine.py

    def relu(self):
        out = Value(0 if self.data < 0 else self.data, (self,), 'ReLU')

        def _backward():
            self.grad += (out.data > 0) * out.grad
        out._backward = _backward

        return out

Прямой проход релу простой: если число отрицательное, результат ноль, иначе результат равен самому числу. Узел out снова получает единственного родителя self и метку 'ReLU'.

Обратите внимание на общий шаблон, который повторяется в каждой операции: взять данные входов, вычислить data результата по обычной математической формуле, создать новый Value с этими данными и списком родителей, прикрепить к нему функцию _backward, вернуть его. Этот шаблон и есть вся суть построения графа вычислений: граф строится сам собой, попутно с обычными вычислениями, вам не нужно отдельно его описывать.

Производные операции: минус, вычитание, деление

Чтобы не писать _backward заново для каждой мыслимой операции, micrograd выражает часть операций через уже определённые:

# Фрагмент: micrograd/micrograd/engine.py

    def __neg__(self): # смена знака: -self
        return self * -1

    def __radd__(self, other): # сложение, когда Value справа: other + self
        return self + other

    def __sub__(self, other): # вычитание: self - other
        return self + (-other)

    def __rsub__(self, other): # вычитание, когда Value справа: other - self
        return other + (-self)

    def __rmul__(self, other): # умножение, когда Value справа: other * self
        return self * other

    def __truediv__(self, other): # деление: self / other
        return self * other**-1

    def __rtruediv__(self, other): # деление, когда Value справа: other / self
        return other * self**-1

Унарный минус __neg__ это просто умножение на -1: новый _backward для этой операции не нужен, вся работа уже проделана внутри __mul__. Вычитание __sub__ выражено как сложение с отрицанием: self - other превращается в self + (-other). Деление __truediv__ выражено через возведение в степень -1: self / other это self * other**-1, то есть умножение на обратную величину. Это изящный приём: вместо того чтобы отдельно продумывать правило дифференцирования для вычитания и деления, достаточно свести их к сложению, умножению и степени, для которых правило уже написано.

Отдельного внимания заслуживают методы с префиксом r: __radd__, __rsub__, __rmul__, __rtruediv__. Они нужны, потому что Python по-разному обрабатывает выражения вида a + b в зависимости от того, с какой стороны стоит объект нужного типа. Когда вы пишете a + b, где a это Value, Python сначала пробует вызвать a.__add__(b), и это работает. Но если вы пишете 2 + a, Python сначала пробует вызвать (2).__add__(a), а у обычного числа int нет метода, который знает, что делать с Value. Тогда Python в качестве запасного варианта пробует вызвать a.__radd__(2), именно его и определяет micrograd. Без __radd__ выражение 2 + a выбросило бы исключение TypeError.

Та же история с __rmul__, она особенно важна на практике: когда мы будем складывать список значений через sum(values), встроенная функция sum по умолчанию начинает с нуля и вычисляет 0 + values[0] + values[1] + .... Самое первое сложение это 0 + values[0], то есть int.__add__ не справится, и в игру вступит values[0].__radd__(0). Без этого метода суммирование списка Value через sum() не работало бы.

Пример: граф для своей формулы

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

# Наш пример
from micrograd.engine import Value

base_fee = Value(50.0)          # базовая плата за доставку
weight = Value(3.0)             # вес посылки, кг
distance = Value(12.0)          # расстояние, км
discount = Value(-8.0)          # сезонная скидка: уменьшает итоговую стоимость

weight_fee = weight * 15.0      # плата за вес: 15 за килограмм
distance_fee = distance * 2.0   # плата за километр

raw_cost = base_fee + weight_fee + distance_fee + discount
cost = raw_cost.relu()          # стоимость не может быть отрицательной

print(cost.data)

Каждая строка с арифметикой на самом деле не просто считает число, а создаёт новый узел Value и связывает его с родителями через _prev. К моменту, когда мы напечатали cost.data (а это 111.0, посчитайте сами: 50 + 45 + 24 - 8), где-то в памяти уже лежит целый граф: cost указывает на raw_cost, тот указывает на четыре слагаемых, а weight_fee и distance_fee в свою очередь указывают на weight, distance и обёрнутые в Value числа 15.0 и 2.0.

Чтобы увидеть этот граф своими глазами, не прибегая к graphviz, напишем простой рекурсивный обход:

# Наш пример
def print_graph(value, depth=0, visited=None):
    """рекурсивно печатает граф вычислений, начиная с заданного узла"""
    if visited is None:
        visited = set()
    indent = '  ' * depth
    op_label = value._op if value._op else 'leaf'  # у листьев операции нет
    print(f"{indent}Value(data={value.data:.2f}) op={op_label}")
    if value in visited:
        return  # узел уже раскрыт выше, не повторяем его детей
    visited.add(value)
    for child in value._prev:
        print_graph(child, depth + 1, visited)

print_graph(cost)

Функция печатает текущий узел с отступом, соответствующим глубине в графе, а затем рекурсивно спускается к каждому из его _prev. Множество visited защищает от повторного раскрытия детей одного и того же узла, если он встречается в графе несколько раз (например, если бы мы использовали weight ещё где-то). Запустив её на cost, вы увидите дерево операций: ReLU на вершине, под ним цепочка +, а в самом низу листья, то есть исходные числа без операции.

Обратите внимание, что _prev это множество, поэтому порядок, в котором родители печатаются при обходе, не гарантирован и может отличаться от запуска к запуску. Для отладки это не страшно: важен сам факт связей, а не порядок их перечисления.

Ограничения micrograd и зачем он всё равно нужен

У micrograd есть сознательные ограничения. Каждый Value хранит одно число, а не массив чисел, поэтому для обучения реальной сети с тысячами параметров пришлось бы создавать тысячи отдельных объектов Value и выполнять операции поштучно в цикле на чистом Python. Это на порядки медленнее, чем векторизованные операции над тензорами в PyTorch или NumPy, где одна инструкция процессора или видеокарты обрабатывает сразу целый массив чисел.

Но именно эта простота и есть ценность micrograd как учебного инструмента. PyTorch устроен по тому же принципу: каждый torch.Tensor с requires_grad=True хранит граф вычислений, узлы которого это тензоры, а не скаляры, а операции между тензорами так же создают новые узлы с прикреплёнными правилами для обратного прохода. Разобравшись на простых скалярных примерах, как работает связка data, grad, родители и операция, вы прочитаете внутренности любого фреймворка глубокого обучения намного увереннее.

Попробуйте сами

Добавьте в класс Value метод exp(self), который в прямом проходе вычисляет out = Value(math.exp(self.data), (self,), 'exp'), то есть экспоненту от текущего значения, с единственным родителем self и меткой 'exp'. Не забудьте добавить import math в начало файла engine.py, иначе math.exp будет недоступен. Обратный проход (_backward) на этом шаге можно оставить пустышкой или, со звёздочкой, попробовать написать его самостоятельно, вспомнив, что производная экспоненты равна самой экспоненте. Критерий успеха: Value(1.0).exp().data близко к 2.71828.
Используя уже существующие операции и метод exp из задания 1, выразите гиперболический тангенс через формулу tanh(x) = (e^(2x) - 1) / (e^(2x) + 1), не трогая внутренности класса Value, то есть напишите обычную функцию my_tanh(x), которая принимает Value и возвращает Value, пользуясь только +, -, *, / и exp(). Критерий успеха: для x = Value(0.0) результат близок к 0.0, а для больших по модулю x результат близок к 1.0 или -1.0.
Придумайте свою формулу с пятью-шестью операциями (сложение, умножение, степень, relu), постройте для неё граф вычислений и воспользуйтесь функцией print_graph из этого урока, чтобы распечатать его. Посчитайте вручную, сколько всего узлов Value получилось, включая листья, и сверьте с тем, что реально печатает обход (подсказка: удобно завести ещё одну функцию, которая просто собирает все уникальные узлы в множество и возвращает его длину).

Итоги

Автоматическое дифференцирование в режиме reverse mode запоминает историю вычислений во время прямого прохода, а затем вычисляет все производные за один обратный проход, в отличие от медленного численного и громоздкого символьного дифференцирования.
Класс Value хранит число data, накопитель градиента grad, множество родителей _prev, метку операции _op и функцию _backward, которая пока остаётся заглушкой до следующего урока.
Каждая операция (__add__, __mul__, __pow__, relu) оборачивает числа в Value, создаёт новый узел с правильными родителями и меткой, попутно пряча в _backward правило для будущего обратного прохода.
Вычитание, деление и унарный минус не требуют отдельной реализации _backward: они выражены через сложение, умножение и степень.
Методы __radd__ и __rmul__ нужны, чтобы выражения вида 2 * a и суммирование через sum() работали, когда Value стоит справа от оператора.
Граф вычислений это DAG из объектов Value, который легко обойти обычной рекурсией по _prev, не прибегая к graphviz, и тот же принцип лежит в основе тензорного автодиффа в PyTorch.

Первоисточники

Первоисточник: видео Андрея Карпати (на английском): The spelled-out intro to neural networks and backpropagation: building micrograd
Код урока на GitHub: karpathy/micrograd/micrograd/engine.py

Видео приведены только как ссылки на первоисточник; текст урока не является их переводом или пересказом. Полные тексты лицензий: в уроке «Что дальше, лицензии и источники».

micrograd: значение и граф вычислений