optim.py

1# this file takes the gradients calculated so far and updates them23import math4import numpy as np56class Adam:7    def __init__(self, params, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.0):8        self.m = {k: np.zeros_like(v) for k, v in params.items()} # running average of the gradient (which direction to go)9        self.v = {k: np.zeros_like(v) for k, v in params.items()} # running average of gradient squared (how big the gradients usually are)10        self.t = 0 # step counter, needed for bias correction11        self.b1, self.b2, self.eps, self.wd = beta1, beta2, eps, weight_decay1213    def step(self, params, grads, lr):14        self.t += 1 # first step is t=1, so 1 - b1**t is never zero1516        for k in params:17            g = grads[k]1819            self.m[k] = self.b1 * self.m[k] + (1 - self.b1) * g # mostly old direction, a little new gradient20            self.v[k] = self.b2 * self.v[k] + (1 - self.b2) * g * g # same idea, but for the squared gradient2122            mhat = self.m[k] / (1 - self.b1 ** self.t) # m starts at 0 so it is too small early on, this scales it back up23            vhat = self.v[k] / (1 - self.b2 ** self.t)2425            update = mhat / (np.sqrt(vhat) + self.eps) # direction divided by typical size, so every number moves about lr per step2627            if params[k].ndim >= 2: # weight decay only for matrices, not biases or layer norm gamma/beta28                update = update + self.wd * params[k] # shrink the weight a little, outside the m/v part (this is the "W" in AdamW)2930            params[k] -= lr * update # changes the array in place, nothing to return3132def lr_schedule(step, warmup, total, lr_max, lr_min):33    if step <= warmup:34        return lr_max * step / warmup # warm-up: ramp from near 0 up to lr_max, a fresh model is fragile3536    frac = (step - warmup) / (total - warmup) # how far through the decay we are, goes from 0 to 137    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * frac)) # cosine slide from lr_max down to lr_min383940def clip_grad_norm(grads, max_norm):41    total = 0.042    for k in grads:43        total += (grads[k] ** 2).sum() # add up the squares of every gradient number in every tensor44    norm = np.sqrt(total) # length of the whole gradient as one big vector4546    if norm > max_norm:47        for k in grads:48            grads[k] *= max_norm / norm # shrink every tensor by the same factor, direction stays, length becomes max_norm4950    return float(norm) # the length before clipping, train.py prints it as |g|