Training framework summary - From Neuron to GPT

Over the previous 16 chapters, we started from scratch and extended the code a little in every chapter, building step by step the tensor, computational graph, network layers, loss functions, optimizers, scheduler, trainer, and more, which are the building blocks for crossing from deep learning to large language models.

Now it is time to collect all these modules and assemble them into a complete, reusable training framework for network models.

This chapter is a summary of the code of the previous 16 chapters, and also the starting point of the next part, large language models. We gather the basic components into one code library, so that later parts no longer have to repeat the code of these components every time.

import contextlib
import math
from abc import ABC, abstractmethod

import numpy as np

DTYPE = np.float32

Tensor¶

The tensor is the cornerstone of the whole framework and the carrier of automatic differentiation. We have extended it with some important features:

Operator overloading: we had already implemented tensor addition, and now we have fully implemented subtraction, multiplication, division, matrix multiplication, and matrix transpose. Tensors can be combined directly with +, -, *, /, and @, just like numbers, and the computational graph is recorded automatically, so they blend into the whole multi-layer model;
Broadcasting support: computation between tensors of different shapes, for example adding a (2, 6) tensor and a (1, 6) tensor. This is very common with complex models and multi-dimensional data;
Disabling gradients (no_grad): when the computational graph is not needed, such as during inference, gradient computation can be switched off temporarily to save memory and compute;
A unified data type (DTYPE = np.float32): the values of tensors are explicitly defined as 32-bit floating-point numbers.

⚠️ Note: Both sides of an operator must be tensors. Operating directly between a tensor and a Python number (for example x * 2) raises an error; wrap the number in a tensor first, for example x * Tensor(2.0).

class Tensor:
    grad_enabled = True

    # Context manager: disable gradients. While the with block runs, the computational graph, backpropagation, and parameter updates are paused
    @classmethod
    @contextlib.contextmanager
    def no_grad(cls):
        prev = cls.grad_enabled
        cls.grad_enabled = False
        try:
            yield
        finally:
            cls.grad_enabled = prev

    def __init__(self, data):
        self.data = np.asarray(data, dtype=DTYPE)
        self.grad = np.zeros_like(self.data)
        self.gradient_fn = None
        self.parents = []

    def backward(self):
        topo = []
        visited = set()
        stack = [(self, False)]

        while stack:
            node, expanded = stack.pop()
            if node in visited:
                continue

            if expanded:
                visited.add(node)
                topo.append(node)
            else:
                stack.append((node, True))
                for p in node.parents:
                    if p not in visited:
                        stack.append((p, False))

        self.grad = np.ones_like(self.data)
        for t in reversed(topo):
            if t.gradient_fn is not None:
                t.gradient_fn()

        for t in topo:
            t.gradient_fn = None
            t.parents = []

    @property
    def shape(self):
        return self.data.shape

    @property
    def dtype(self):
        return self.data.dtype

    # Tensor addition
    def __add__(self, other):
        p = Tensor(self.data + other.data)

        def gradient_fn():
            self.grad += self._unbroadcast(p.grad, self.shape)
            other.grad += self._unbroadcast(p.grad, other.shape)

        return p.attach(gradient_fn, parents=[self, other])

    # Tensor subtraction
    def __sub__(self, other):
        p = Tensor(self.data - other.data)

        def gradient_fn():
            self.grad += self._unbroadcast(p.grad, self.shape)
            other.grad += self._unbroadcast(-p.grad, other.shape)

        return p.attach(gradient_fn, parents=[self, other])

    # Tensor multiplication
    def __mul__(self, other):
        p = Tensor(self.data * other.data)

        def gradient_fn():
            self.grad += self._unbroadcast(p.grad * other.data, self.shape)
            other.grad += self._unbroadcast(p.grad * self.data, other.shape)

        return p.attach(gradient_fn, parents=[self, other])

    # Tensor division
    def __truediv__(self, other):
        p = Tensor(self.data / other.data)

        def gradient_fn():
            self.grad += self._unbroadcast(p.grad / other.data, self.shape)
            other.grad += self._unbroadcast(-p.grad * self.data / (other.data ** 2), other.shape)

        return p.attach(gradient_fn, parents=[self, other])

    # Tensor matrix multiplication (@)
    def __matmul__(self, other):
        p = Tensor(np.matmul(self.data, other.data))

        def gradient_fn():
            self.grad += self._unbroadcast(np.matmul(p.grad, other.data.swapaxes(-1, -2)), self.shape)
            other.grad += self._unbroadcast(np.matmul(self.data.swapaxes(-1, -2), p.grad), other.shape)

        return p.attach(gradient_fn, parents=[self, other])

    # Tensor matrix transpose
    def transpose(self, axes=None):
        p = Tensor(np.transpose(self.data, axes))

        def gradient_fn():
            if axes is None:
                self.grad += np.transpose(p.grad)
            else:
                idx = np.argsort(axes)
                self.grad += np.transpose(p.grad, idx)

        return p.attach(gradient_fn, parents=[self])

    # Tensor matrix transpose
    @property
    def T(self):
        return self.transpose()

    # Tensor reshape
    def reshape(self, shape):
        p = Tensor(np.reshape(self.data, shape))

        def gradient_fn():
            self.grad += np.reshape(p.grad, self.shape)

        return p.attach(gradient_fn, parents=[self])

    # Set the gradient (closure) function and parent list of the tensor
    def attach(self, gradient_fn, parents):
        if Tensor.grad_enabled:
            self.gradient_fn = gradient_fn
            self.parents = parents
        return self

    def __repr__(self):
        return f"Tensor(shape={self.shape}, dtype={self.dtype})"

    # Broadcasting
    @staticmethod
    def _unbroadcast(grad, shape):
        if grad.ndim > len(shape):
            grad = grad.sum(axis=tuple(range(grad.ndim - len(shape))))

        for axis, dim in enumerate(shape):
            if dim == 1 and grad.shape[axis] != 1:
                grad = grad.sum(axis=axis, keepdims=True)
        return grad.reshape(shape)

Model¶

A neural network model is made up of multiple network layers. All layers support the basic functions defined by the abstract layer.

Abstract layer¶

The abstract layer defines the most important method interfaces:

forward (inference method): implements forward propagation;
parameters (parameter list): provides all the parameters that need training;
train (training method): enters training mode;
eval (test method): enters evaluation and inference mode.

class Layer(ABC):

    def __init__(self):
        self.training = True

    def __call__(self, *args):
        return self.forward(*args)

    # Training mode
    def train(self):
        self.training = True

    # Inference mode
    def eval(self):
        self.training = False

    @abstractmethod
    def forward(self, *args):
        pass

    @property
    def parameters(self):
        return []

    def __repr__(self):
        return f"{type(self).__name__}[]"

Linear layer¶

Weighted sum, the most basic network layer.

class Linear(Layer):

    def __init__(self, in_size, out_size):
        super().__init__()
        self.weight = Tensor(np.random.randn(out_size, in_size).astype(DTYPE) * np.sqrt(2 / in_size))
        self.bias = Tensor(np.zeros(out_size).astype(DTYPE))

    def forward(self, x: Tensor):
        p = Tensor(x.data @ self.weight.data.T + self.bias.data)

        def gradient_fn():
            grad = p.grad.reshape(-1, p.grad.shape[-1])
            self.weight.grad += grad.T @ x.data.reshape(-1, x.shape[-1])
            self.bias.grad += np.sum(grad, axis=0)
            x.grad += p.grad @ self.weight.data

        return p.attach(gradient_fn, {x})

    @property
    def parameters(self):
        return [self.weight, self.bias]

Composite layer¶

An abstract layer that can serve as a container for sublayers, and the foundation of complex network layers.

class Composite(Layer, ABC):

    def __init__(self, layers):
        super().__init__()
        self.layers = list(layers)

    def train(self):
        super().train()
        for l in self.layers:
            l.train()

    def eval(self):
        super().eval()
        for l in self.layers:
            l.eval()

    @property
    def parameters(self):
        return [p for l in self.layers for p in l.parameters]

Sequential layer¶

The most basic container of network layers, in which all the sublayers run in order.

class Sequential(Composite):

    def forward(self, x: Tensor):
        for l in self.layers:
            x = l(x)
        return x

Embedding layer¶

Converts the index encodings of characters or words into dense embedding vectors by table lookup, avoiding the sparse vectors of one-hot encoding.

class Embedding(Layer):

    def __init__(self, vocab_size, embedding_size, std=0.02):
        super().__init__()
        self.weight = Tensor(np.random.randn(vocab_size, embedding_size).astype(DTYPE) * std)

    def forward(self, x: Tensor):
        p = Tensor(self.weight.data[x.data.astype(np.int64)])

        def gradient_fn():
            np.add.at(self.weight.grad, x.data.astype(np.int64), p.grad)

        return p.attach(gradient_fn, parents=[self.weight])

    @property
    def parameters(self):
        return [self.weight]

Mean pooling layer¶

Merges a group of vectors in a sequence into one vector by taking their mean.

class MeanPool(Layer):

    def forward(self, x: Tensor):
        p = Tensor(np.mean(x.data, axis=1))

        def gradient_fn():
            x.grad += p.grad[:, None, :] / x.data.shape[1]

        return p.attach(gradient_fn, parents=[x])

Activation functions¶

Activation functions introduce nonlinearity into the network model, so that the model can learn complex relationships.

Rectified linear unit (ReLU)¶

The default first choice for hidden layers. It is fast to compute and does not cause vanishing gradients.

class ReLU(Layer):

    def forward(self, x: Tensor):
        a = Tensor(np.maximum(0, x.data))

        def gradient_fn():
            x.grad += a.grad * (a.data > 0)

        return a.attach(gradient_fn, parents=[x])

Hyperbolic tangent (Tanh)¶

The output is centered on 0 and limited to the range (-1, 1). It suffers from gradient saturation.

class Tanh(Layer):

    def forward(self, x: Tensor):
        a = Tensor(np.tanh(x.data))

        def gradient_fn():
            x.grad += a.grad * (1 - a.data ** 2)

        return a.attach(gradient_fn, parents=[x])

Sigmoid¶

Used in the output layer for binary classification problems.

class Sigmoid(Layer):

    def __init__(self, clip_range=(-100, 100)):
        super().__init__()
        self.clip_range = clip_range

    def forward(self, x: Tensor):
        z = np.clip(x.data, self.clip_range[0], self.clip_range[1])
        a = Tensor(1 / (1 + np.exp(-z)))

        def gradient_fn():
            x.grad += a.grad * a.data * (1 - a.data)

        return a.attach(gradient_fn, parents=[x])

Softmax¶

Used in the output layer for multi-class classification problems.

class Softmax(Layer):

    def __init__(self, axis=-1):
        super().__init__()
        self.axis = axis

    def forward(self, x: Tensor):
        exp = np.exp(x.data - np.max(x.data, axis=self.axis, keepdims=True))
        a = Tensor(exp / np.sum(exp, axis=self.axis, keepdims=True))

        def gradient_fn():
            grad = np.sum(a.data * a.grad, axis=self.axis, keepdims=True)
            x.grad += a.data * (a.grad - grad)

        return a.attach(gradient_fn, parents=[x])

Loss functions¶

A loss function measures the error between the prediction and the label. It is the “baton” of training and the starting point of backpropagation.

Abstract class¶

The abstract class of loss functions defines a single method interface:

loss (loss method): computes the loss.

class Loss(ABC):

    def __call__(self, *args, **kwargs):
        return self.loss(*args, **kwargs)

    @abstractmethod
    def loss(self, *args, **kwargs):
        pass

Mean squared error loss function¶

Used for regression problems.

class MSELoss(Loss):

    def loss(self, p: Tensor, y: Tensor):
        mse = Tensor(np.mean(np.square(y.data - p.data)))

        def gradient_fn():
            p.grad += mse.grad * (-2 * (y.data - p.data) / y.data.size)

        return mse.attach(gradient_fn, parents=[p])

Cross-entropy loss function¶

Used for multi-class classification problems. It already includes the Softmax computation internally.

class CELoss(Loss):

    def loss(self, p: Tensor, y: Tensor):
        exp = np.exp(p.data - np.max(p.data, axis=-1, keepdims=True))
        softmax = exp / np.sum(exp, axis=-1, keepdims=True)

        if y.data.ndim == softmax.ndim:
            labels = y.data
        else:
            labels = np.zeros_like(softmax)
            rows = np.arange(y.data.size)
            labels[rows, y.data.astype(np.int64)] = 1

        log = np.log(np.clip(softmax, 1e-10, 1.0))
        ce = Tensor(-np.mean(np.sum(labels * log, axis=-1)))

        def gradient_fn():
            grad = (softmax - labels) / labels.shape[0]
            p.grad += ce.grad * grad

        return ce.attach(gradient_fn, parents=[p])

Binary cross-entropy loss function¶

Used for binary classification problems.

class BCELoss(Loss):

    def loss(self, p: Tensor, y: Tensor):
        clipped = np.clip(p.data, 1e-7, 1 - 1e-7)
        bce = Tensor(-np.mean(y.data * np.log(clipped) + (1 - y.data) * np.log(1 - clipped)))

        def gradient_fn():
            grad = (clipped - y.data) / (clipped * (1 - clipped)) / y.data.size
            p.grad += bce.grad * grad

        return bce.attach(gradient_fn, parents=[p])

Optimizers¶

An optimizer updates the parameters according to the gradients and the learning rate.

Abstract class¶

The abstract class of optimizers defines two basic method interfaces:

step (step method): updates the parameters;
zero_grad (zero-gradient method): clears the gradients of the parameters before each iteration.

class Optimizer(ABC):

    def __init__(self, parameters, lr):
        self.parameters = list(parameters)
        self.lr = lr

    @abstractmethod
    def step(self):
        pass

    def zero_grad(self):
        for p in self.parameters:
            p.grad = np.zeros_like(p.data)

Stochastic gradient descent optimizer¶

The simplest and fastest, with a fixed learning rate.

class SGDOptimizer(Optimizer):

    def step(self):
        for p in self.parameters:
            p.data -= p.grad * self.lr

Adaptive moment estimation optimizer¶

Adjusts the effective step size of each parameter individually, and comes with gradient clipping.

class AdamOptimizer(Optimizer):

    def __init__(self, parameters, lr=0.01, betas=(0.9, 0.999), eps=1e-8):
        super().__init__(parameters, lr)
        self.b1, self.b2 = betas
        self.eps = eps
        self.m = [None] * len(parameters)
        self.v = [None] * len(parameters)
        self.t = 0

    def step(self):
        self.t += 1
        for i, p in enumerate(self.parameters):
            if p is not None:
                if self.m[i] is None:
                    self.m[i] = np.zeros_like(p.data)
                    self.v[i] = np.zeros_like(p.data)

                self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * p.grad
                self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * (p.grad ** 2)
                m = self.m[i] / (1 - self.b1 ** self.t)
                v = self.v[i] / (1 - self.b2 ** self.t)
                self._apply_weight_decay(p)
                p.data -= self.lr * m / (np.sqrt(v) + self.eps)

    def _apply_weight_decay(self, p):
        pass

    def clip_grad_norm(self, max_norm=1.0):
        sq = 0.0
        for p in self.parameters:
            sq += float(np.sum(p.grad ** 2))

        if np.sqrt(sq) > max_norm > 0:
            scale = max_norm / (np.sqrt(sq) + 1e-6)
            for p in self.parameters:
                p.grad *= scale

Adaptive moment estimation with weight decay optimizer¶

Adam with weight decay.

class AdamWOptimizer(AdamOptimizer):

    def __init__(self, parameters, lr=0.01, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01):
        super().__init__(parameters, lr, betas, eps)
        self.weight_decay = weight_decay

    def _apply_weight_decay(self, p):
        if p.data.ndim >= 2:
            p.data -= p.data * self.weight_decay * self.lr

Learning rate scheduler¶

A learning rate scheduler adjusts the learning rate dynamically during training (automatic tuning).

Abstract class¶

The abstract class of learning rate schedulers defines a single method interface:

step (step method): adjusts the learning rate according to the current step count.

class LRScheduler(ABC):

    @abstractmethod
    def step(self, current_step):
        pass

Warmup cosine annealing scheduler¶

Lets the learning rate climb from 0 (warmup) and then fall smoothly along a cosine curve (cosine annealing). This is the standard practice for training large models.

class WarmupCosineScheduler(LRScheduler):

    def __init__(self, lr, total_steps, warmup_steps, min_lr=0.0):
        self.lr = lr
        self.total_steps = max(total_steps, 1)
        self.warmup_steps = max(min(warmup_steps, self.total_steps), 0)
        self.min_lr = min_lr

    def step(self, current_step):
        if self.warmup_steps > 0 and current_step < self.warmup_steps:
            return self.lr * (current_step + 1) / self.warmup_steps

        if current_step >= self.total_steps:
            return self.min_lr

        progress = (current_step - self.warmup_steps) / max(self.total_steps - self.warmup_steps, 1)
        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
        return self.min_lr + (self.lr - self.min_lr) * cosine

Dataset¶

A dataset loads the data, splits it into training and test sets, and walks through all the samples in batches.

Abstract class¶

The abstract class of datasets defines the basic method interfaces:

load (load method): loads the data;
train (training method): provides the training set data;
eval (test method): provides the test set data;
all (all samples): reads all the samples;
__len__ (number of samples): returns the number of samples;
__getitem__ (a single sample): reads a sample by index.

class Dataset(ABC):

    def __init__(self, batch_size=1):
        self.batch_size = batch_size
        self.load()
        self.train()

    @abstractmethod
    def load(self):
        pass

    def train(self):
        self.data = self.train_data

    def eval(self):
        self.data = self.test_data

    def all(self):
        x, y = self.data
        return Tensor(x), Tensor(y)

    def __len__(self):
        x, *_ = self.data
        return len(x) // self.batch_size

    def __getitem__(self, index):
        s = slice(index * self.batch_size, (index + 1) * self.batch_size)
        x, y = self.data
        return Tensor(x[s]), Tensor(y[s])

Trainer¶

A trainer packages the model, the loss function, and the optimizer together, and implements the training and testing logic.

Abstract class¶

The abstract class of trainers defines two basic method interfaces:

train (training method): trains for multiple epochs on the training set of a dataset, and can be connected to a learning rate scheduler;
test (test method): evaluates the model on the test set of a dataset.

class Trainer(ABC):

    def __init__(self, model, loss_fn, optimizer):
        self.model = model
        self.loss_fn = loss_fn
        self.optimizer = optimizer

    @abstractmethod
    def train(self, dataset, epochs, scheduler=None):
        pass

    @abstractmethod
    def test(self, dataset):
        pass

With this, our core framework has basically taken shape.

Although this framework is written with only NumPy, it has the core features of a modern deep learning framework. In the next part, we will build on it to construct a real GPT model.

Exercises¶

Tensors currently do not support direct operations with Python numbers. Modify __add__, __sub__, __mul__, and __truediv__ so that x * 2 and x + 1.0 also work, and the gradients in backpropagation remain correct. Hint: at the start of each operation, wrap the number in a tensor first.