Un autograd desde cero

Deep learning a secas. Episodio seis: un autograd desde cero. En el episodio cinco calculamos los gradientes a mano, con la regla de la cadena. Hoy escribimos el programa que lo hace solo, en unas cuarenta líneas de Python, y lo comparamos contra PyTorch.

Qué vamos a construir

• Un número que recuerda de dónde salió
• Cada operación sabe su derivada local
• backward recorre el grafo de atrás para adelante
• Mismo resultado que torch.autograd

Autograd viene de automatic gradient. Es lo que corre cuando llamás a loss punto backward, como en Alineación, episodio dos. Cuatro pasos. Primero, un número que además de su valor recuerda de qué operación salió. Segundo, cada operación sabe una sola cosa: su derivada local, la de ella sola. Tercero, una función backward que recorre todas las operaciones en orden inverso y encadena esas derivadas. Y cuarto, lo verificamos contra torch punto autograd. La versión que armamos sigue la idea de micrograd, el autograd mínimo de Andrej Karpathy.

Parte 1 · El grafo de cómputo

Para derivar una cuenta hay que saber cómo se hizo. Así que primero, que la cuenta se anote a sí misma.

Un número que se acuerda de dónde viene

La pieza es una clase, Value: un número envuelto. Guarda su valor en data, y un grad que arranca en cero. grad va a ser la derivada de la pérdida final respecto de este número. Además guarda de qué Values salió, sus hijos, y con qué operación. El truco está en los métodos con doble guion bajo. Cuando Python ve a por b con dos Values, llama a a punto mul de b. Devolvemos un Value nuevo que anota a sus dos padres: operator overloading. Así, escribir e igual a por b arma el grafo sin que nadie lo pida. e vale seis, con la operación por y sus hijos dos y tres. d vale dos: seis más menos cuatro.

class Value:
    def __init__(self, data, hijos=(), op=""):
        self.data = data          # el número
        self.grad = 0.0           # dL/d(este número), lo llena backward
        self._prev = tuple(hijos) # de qué Values salió
        self._op = op             # con qué operación
    def __add__(self, otro):      # a + b
        return Value(self.data + otro.data, (self, otro), "+")
    def __mul__(self, otro):      # a * b
        return Value(self.data * otro.data, (self, otro), "*")

a, b, c = Value(2.0), Value(3.0), Value(-4.0)
e = a * b                         # Python llama a.__mul__(b)
d = e + c
print("e =", e.data, " op:", e._op, " hijos:", [h.data for h in e._prev])
print("d =", d.data, " op:", d._op, " hijos:", [h.data for h in d._prev])

e = 6.0  op: *  hijos: [2.0, 3.0]
d = 2.0  op: +  hijos: [6.0, -4.0]

La cuenta, dibujada como grafo

Esta es la cuenta del episodio: L igual a, entre paréntesis, a por b más c, todo por f. Las cajas son Values y los círculos, operaciones. Cada flecha es un hijo anotado en underscore prev. Eso es un grafo de cómputo. El forward: los valores corren de izquierda a derecha, en el dibujo, los puntos ámbar viajando por las flechas. a dos por b tres da e seis. Más c, menos cuatro, da d dos. Por f, menos dos, da L, menos cuatro. Todavía no hay ninguna derivada: solo quedó registrado el camino.

Parte 2 · Derivadas locales

Cada operación sabe derivarse a sí misma, y nada más. El resto lo hace la regla de la cadena.

Lo único que sabe cada operación

la suma: mover e en 1 mueve d en 1

el producto: la derivada respecto de uno es el otro

regla de la cadena: lo que llega de arriba × la local

La derivada parcial, esa d redondeada, mide cuánto cambia una cosa si movés otra un poquito y dejás quieto lo demás. La suma. Si e crece uno, d crece uno: derivada uno, para los dos. En el dibujo, una neurona con los dos pesos en uno. El producto. e igual a a por b: con b fija en tres, e es tres veces a. Cada unidad de a suma tres. La derivada respecto de a es b, y respecto de b es a. En el dibujo, a entra con peso tres. Y la regla de la cadena del episodio cinco: la derivada de L respecto de a es lo que llega desde arriba, la derivada de L respecto de e, por la derivada local. Menos dos por tres, menos seis. En el dibujo, el gradiente entra y se multiplica por la local.

La derivada local, guardada en una función

Ahora en código. Adentro de mul definimos una función chiquita, underscore backward. Es una closure: una función que se acuerda de self, otro y out. Toma el gradiente que llegó a la salida, out punto grad, y lo multiplica por la derivada local: a self le toca el valor de otro, y a otro, el de self. Fijate el más igual: suma, no pisa. La función se guarda en el Value de salida, sin ejecutarse. La suma es igual de simple: copia el gradiente de la salida a los dos hijos, porque la derivada local es uno. Para probar, ponemos a mano el grad de e en uno y llamamos a su backward. a recibe tres, el valor de b, y b recibe dos, el valor de a.

class Value:
    def __init__(self, data, hijos=(), op=""):
        self.data, self.grad = data, 0.0
        self._prev, self._op = tuple(hijos), op
        self._backward = lambda: None
# --- desde acá, lo que se ve en el video ---
    def __mul__(self, otro):
        out = Value(self.data * otro.data, (self, otro), "*")
        def _backward():                      # regla local de *
            self.grad += otro.data * out.grad     # d(ab)/da = b
            otro.grad += self.data * out.grad     # d(ab)/db = a
        out._backward = _backward             # se guarda, no se corre
        return out
    def __add__(self, otro):
        out = Value(self.data + otro.data, (self, otro), "+")
        def _backward():                      # regla local de +
            self.grad += out.grad                 # d(a+b)/da = 1
            otro.grad += out.grad                 # d(a+b)/db = 1
        out._backward = _backward
        return out
a, b = Value(2.0), Value(3.0)
e = a * b
e.grad = 1.0          # dL/de, supongamos que vale 1
e._backward()         # reparte hacia a y b
print("a.grad =", a.grad, "  b.grad =", b.grad)

a.grad = 3.0   b.grad = 2.0

Parte 3 · backward, en orden

Cada nodo sabe repartir su gradiente. Falta llamarlos en el orden correcto: de la salida hacia las hojas.

El orden topológico

Un nodo reparte su gradiente cuando ya recibió todo lo de arriba. Para eso armamos un orden topológico: una lista donde cada nodo aparece después de todos sus hijos. visitar baja primero por los hijos, de forma recursiva. Y recién cuando terminó con ellos, agrega el nodo a la lista. El gradiente de L respecto de sí misma es uno: es el punto de partida. Y recorremos la lista al revés, de la salida a las hojas, llamando al backward de cada uno. El orden queda: dos, tres, el por, menos cuatro, el más, menos dos, y el por final. Al revés, L va primero.

def backward(raiz):
    orden, vistos = [], set()
    def visitar(v):                 # orden topológico:
        if v not in vistos:         # cada nodo va después
            vistos.add(v)           # de todos sus hijos
            for hijo in v._prev:
                visitar(hijo)
            orden.append(v)
    visitar(raiz)
    raiz.grad = 1.0                 # dL/dL = 1
    for v in reversed(orden):       # de la salida a las hojas
        v._backward()
    return orden

a, b, c, f = Value(2.0), Value(3.0), Value(-4.0), Value(-2.0)
L = (a * b + c) * f
orden = backward(L)
print("orden:", " ".join(str(v._op or v.data) for v in orden))

orden: 2.0 3.0 * -4.0 + -2.0 *

El backward recorre la lista al revés

Los siete nodos en el orden de la lista: hojas primero, salida al final. El backward arranca por el último, L, con gradiente uno. Su operación es un producto, así que le pasa a d el valor de f y a f el valor de d. Ahora d tiene todo su gradiente y lo puede pasar. Es una suma: lo copia a e y a c. e es un producto: le pasa a a y a b. Y llegamos a las hojas, que no tienen nada que repartir.

a = 2 → b = 3 → e = a·b → c = −4 → d = e + c → f = −2 → L = d·f

El ejemplo completo

Todo junto. Las mismas cuentas de antes, con Values. L da menos cuatro. Una sola llamada a backward. Y recorremos los siete nodos imprimiendo valor y gradiente. L tiene gradiente uno. d, menos dos, que es el valor de f. f, dos, que es el valor de d. e y c, menos dos, copiados por la suma. a, menos seis. Y b, menos cuatro. Si subís a un poquito, L baja seis veces eso.

a, b = Value(2.0), Value(3.0)
c, f = Value(-4.0), Value(-2.0)
e = a * b          # 6
d = e + c          # 2
L = d * f          # -4
backward(L)
for nombre, v in [("L", L), ("d", d), ("f", f), ("e", e),
                  ("c", c), ("a", a), ("b", b)]:
    print(f"{nombre}: valor={v.data:5}  grad={v.grad:5}")

L: valor= -4.0  grad=  1.0
d: valor=  2.0  grad= -2.0
f: valor= -2.0  grad=  2.0
e: valor=  6.0  grad= -2.0
c: valor= -4.0  grad= -2.0
a: valor=  2.0  grad= -6.0
b: valor=  3.0  grad= -4.0

Forward y backward sobre el mismo grafo

El mismo grafo, completo. Las cajas y las operaciones. Arriba de cada caja, en ámbar, el valor del forward, que viaja de izquierda a derecha. Y abajo, en rojo, el gradiente del backward, que viaja al revés: los puntos rojos van de L hacia las hojas. Arranca en uno; d recibe menos dos y f, dos; la suma copia menos dos a e y a c; y el primer producto da menos seis para a y menos cuatro para b. El recorrido del episodio cinco, hecho por un bucle.

Lo que hizo backward, paso a paso

L = d·f: el producto final

la suma copia el gradiente

el primer producto: a recibe b, b recibe a

Las mismas tres etapas, con fórmulas. El producto final. La derivada de L respecto de d es uno, lo que llega, por f: menos dos. Y respecto de f, uno por d: dos. En el dibujo entra el uno y sale menos dos. La suma copia: e y c reciben menos dos por uno, menos dos. El primer producto: a recibe menos dos por b, que es tres, y da menos seis. b recibe menos dos por a, que es dos, y da menos cuatro. En el dibujo, la cuenta de b. Cada paso usa el gradiente que llegó y un valor que dejó el forward: por eso el forward va antes.

Parte 4 · ¿Está bien? Comprobarlo

Un autograd que da números equivocados es peor que ninguno. Dos formas de verificarlo: con una diferencia y con PyTorch.

Gradient check, y un número usado tres veces

La primera verificación no necesita ninguna derivada. Movemos a en un paso chiquito, h, diez a la menos cuatro. Y medimos cuánto cambió L, dividido por h. Esa pendiente aproximada se llama gradient check. La segunda prueba es una trampa: x por x más x. Acá x aparece tres veces. Corremos backward sobre y. El autograd y la diferencia coinciden: menos seis. Y la derivada de y en x igual a tres da siete, que es dos x más uno. Si underscore backward usara igual en vez de más igual, cada aporte pisaría al anterior y daría tres: lo probamos, y da tres.

a = Value(2.0)
L = (a * Value(3.0) + Value(-4.0)) * Value(-2.0)
backward(L)

h = 1e-4                                   # un paso chiquito
pendiente = (L_de(2.0 + h) - L_de(2.0)) / h  # cuánto cambió L
print("autograd: dL/da =", a.grad)
print("numérico: dL/da =", round(pendiente, 6))

x = Value(3.0)
y = x * x + x                  # x se usa tres veces
backward(y)
print("x*x + x en x=3: dy/dx =", x.grad, "(2x+1 = 7)")

autograd: dL/da = -6.0
numérico: dL/da = -6.0
x*x + x en x=3: dy/dx = 7.0 (2x+1 = 7)

Las dos comprobaciones

la pendiente, medida de verdad

un aporte por cada uso, sumados: por eso +=

La primera. La derivada es el límite de esta división cuando h se achica: cuánto cambió L, sobre cuánto movimos a. En el dibujo, la recta tangente: su inclinación es la derivada. Es lenta, pero sirve para controlar. La segunda. Cuando una variable se usa en varios lugares, su gradiente es la suma de los aportes de cada camino. En x por x, la primera x recibe el valor de la segunda, tres, y la segunda recibe tres. El más x aporta uno. En el dibujo, las tres barras: tres más tres más uno, siete. Por eso todos los underscore backward suman con más igual.

Una operación nueva: ReLU

Agregar una operación son siempre las mismas dos cosas. Por ejemplo la ReLU, rectified linear unit, del episodio uno. El forward: el máximo entre cero y el valor. La derivada local: uno si la ReLU dejó pasar el valor, cero si lo cortó. Se multiplica por el gradiente que llega. La probamos con un valor positivo y uno negativo. Con dos coma cinco, la salida es dos coma cinco y la derivada, uno. Con menos uno, sale cero y la derivada es cero: por una ReLU apagada no pasa gradiente.

def relu(v):
    out = Value(max(0.0, v.data), (v,), "relu")   # forward
    def _backward():                               # derivada local:
        v.grad += (out.data > 0) * out.grad        # 1 si pasó, 0 si no
    out._backward = _backward
    return out

for z0 in (2.5, -1.0):
    z = Value(z0)
    y = relu(z)
    backward(y)
    print(f"z = {z0:4}  relu = {y.data}  dy/dz = {z.grad}")

z =  2.5  relu = 2.5  dy/dz = 1.0
z = -1.0  relu = 0.0  dy/dz = 0.0

Contra torch.autograd

La prueba de fuego. Una neurona entera con nuestros Values: dos entradas, dos pesos, un sesgo y la ReLU. La misma neurona en PyTorch, con los cinco números en un tensor. requires_grad igual a true le pide que anote las operaciones, lo mismo que hace nuestro Value. Y su backward. Imprimimos sus gradientes en el mismo orden. Coinciden uno por uno: uno, menos dos y uno para los pesos y el sesgo, cero coma cinco y menos uno coma cinco para las entradas. Y la salida, tres coma setenta y cinco, en los dos. El gradiente de cada peso es su entrada; el de cada entrada, su peso.

# nuestra versión: una neurona con ReLU
x1, x2 = Value(1.0), Value(-2.0)
w1, w2, b = Value(0.5), Value(-1.5), Value(0.25)
y = relu(w1 * x1 + w2 * x2 + b)
backward(y)
print("nuestro:", [v.grad for v in (w1, w2, b, x1, x2)])

# PyTorch: los mismos números, requires_grad=True
T = torch.tensor([1.0, -2.0, 0.5, -1.5, 0.25], requires_grad=True)
tx1, tx2, tw1, tw2, tb = T
yt = torch.relu(tw1 * tx1 + tw2 * tx2 + tb)
yt.backward()
print("torch  :", T.grad[[2, 3, 4, 0, 1]].tolist())
print("y:", y.data, "y torch:", yt.item())

nuestro: [1.0, -2.0, 1.0, 0.5, -1.5]
torch  : [1.0, -2.0, 1.0, 0.5, -1.5]
y: 3.75 y torch: 3.75

Paszke et al., 2019 — PyTorch: An Imperative Style, High-Performance Deep Learning Library · arXiv:1912.01703

reverse-mode · «PyTorch uses the operator overloading approach, which builds up a representation of the computed function every time it is executed.»

No es una simplificación: es el mismo diseño. El paper de PyTorch, de dos mil diecinueve, lo dice así: usa operator overloading, que arma una representación de la función calculada cada vez que se ejecuta. Y hace reverse-mode automatic differentiation: el gradiente de una salida escalar respecto de muchas entradas, de atrás para adelante. Eso es nuestro backward. Como el grafo se rearma en cada ejecución, puede tener ifs y bucles de Python adentro.

Nuestro Value y PyTorch

La diferencia es de escala, no de idea. Nuestro Value guarda un número por nodo, tiene tres operaciones con su derivada local, recorre el orden topológico al revés y acumula con más igual. PyTorch guarda tensores enteros por nodo, así que un producto de matrices es un solo nodo; tiene cientos de operaciones con su derivada escrita; y hace el mismo recorrido inverso. También acumula: si llamás dos veces a backward, los gradientes se suman. Por eso el training loop borra los gradientes antes de cada paso. Eso es el episodio siete.

Resumiendo

• Value: un número que anota sus hijos y su operación
• Cada operación guarda su derivada local en _backward
• backward: orden topológico, al revés, con $\partial L/\partial L = 1$
• Un número usado varias veces suma sus aportes: +=
• Gradient check y torch.autograd dan lo mismo
• Próximo: el training loop, con dígitos escritos a mano

Papers del episodio

• Paszke et al., 2019 — PyTorch: An Imperative Style, High-Performance Deep Learning Library