Implementing a straightforward neural network framework from scratch
Despite doing a little work and research within the AI ecosystem for a while, I didn’t truly stop to take into consideration backpropagation and gradient updates inside neural networks until recently. This text seeks to rectify that and can hopefully provide an intensive yet easy-to-follow dive into the subject by implementing a straightforward (yet somewhat powerful) neural network framework from scratch.
Fundamentally, a neural network is only a mathematical function from our input space to our desired output space. In actual fact, we will effectively “unwrap” any neural network right into a function. Consider, as an illustration, the next easy neural network with two layers and one input:
We will now construct an equivalent function by going forwards layer by layer, ranging from the input. Let’s follow our final function layer by layer:
- On the input, we start with the identity function pred(x) = x
- At the primary linear layer, we get pred(x) = w₁x + b₁
- The ReLU nets us pred(x) = max(0, w₁x + b₁)
- At the ultimate layer, we get pred(x) = w₂(max(0, w₁x + b₁)) + b₂
With more complicated nets, these functions in fact get unwieldy, but the purpose is that we will construct such representations of neural networks.
We will go one step further though — functions of this kind are usually not extremely convenient for computation, but we will parse them right into a more useful form, namely a syntax tree. For our easy net, the tree would seem like this:
On this tree form, our leaves are parameters, constants, and inputs, and the opposite nodes are elementary operations whose arguments are their children. In fact, these elementary operations don’t need to be binary — the sigmoid operation, as an illustration, is unary (and so is ReLU if we don’t represent it as a max of 0 and x), and we will decide to support multiplication and addition of multiple input.
By pondering of our network as a tree of those elementary operations, we will now do numerous things very easily with recursion, which is able to form the idea of each our backpropagation and forward propagation algorithms. In code, we will define a recursive neural network class that appears like this:
from dataclasses import dataclass, field
from typing import List@dataclass
class NeuralNetNode:
"""A node in our neural network tree"""
children: List['NeuralNetNode'] = field(default_factory=list)
def op(self, x: List[float]) -> float:
"""The operation that this node performs"""
raise NotImplementedError
def forward(self) -> float:
"""Evaluate this node on the given input"""
return self.op([child.forward() for child in self.children])
# That is only for convenience
def __call__(self) -> List[float]:
return self.forward()
def __repr__(self):
return f'{self.__class__.__name__}({self.children})'
Suppose now that we now have a differentiable loss function for our neural network, say MSE. Recall that MSE (for one sample) is defined as follows:
We now want to update our parameters (the green circles in our tree representation) given the worth of our loss. To do that, we want the derivative of our loss function with respect to every parameter. Calculating this directly from the loss is amazingly difficult though — in spite of everything, our MSE is calculated by way of the worth predicted by our neural net, which might be a very complicated function.
That is where very useful piece of mathematics — the chain rule — comes into play. As an alternative of being forced to compute our highly complex derivatives from the get-go, we will as an alternative compute a series of simpler derivatives.
It seems that the chain rule meshes thoroughly with our recursive tree structure. The thought principally works as follows: assuming that we now have easy enough elementary operations, each elementary operation knows its derivative with respect to all of its arguments. Given the derivative from the parent operation, we will thus compute the derivative of every child operation with respect to the loss function through easy multiplication. For a straightforward linear regression model using MSE, we will diagram it as follows:
In fact, a few of our nodes don’t do anything with their derivatives — namely, only our leaf nodes care. But now each node can get the derivative of its output with respect to the loss function through this recursive process. We will thus add the next methods to our NeuralNetNode class:
def grad(self) -> List[float]:
"""The gradient of this node with respect to its inputs"""
raise NotImplementedErrordef backward(self, derivative_from_parent: float):
"""Propagate the derivative from the parent to the kids"""
self.on_backward(derivative_from_parent)
deriv_wrt_children = self.grad()
for child, derivative_wrt_child in zip(self.children, deriv_wrt_children):
child.backward(derivative_from_parent * derivative_wrt_child)
def on_backward(self, derivative_from_parent: float):
"""Hook for subclasses to override. Things like updating parameters"""
pass
Exercise 1: Try creating considered one of these trees for a straightforward linear regression model and perform the recursive gradient updates by hand for a few steps.
Note: For simplicity’s sake, we require our nodes to have just one parent (or none in any respect). If each node is allowed to have multiple parents, our backwards() algorithm becomes somewhat more complicated as each child must sum the derivative of its parents to compute its own. We will do that iteratively with a topological sort (e.g. see here) or still recursively, i.e. with reverse accumulation (though on this case we would want to do a second pass to truly update all the parameters). This isn’t extraordinarily difficult, so I’ll leave it as an exercise to the reader (and can speak about it more partially 2, stay tuned).
Constructing Models
The remainder of our code really just involves implementing parameters, inputs, and operations, and naturally running our training. Parameters and inputs are fairly easy constructs:
import random@dataclass
class Input(NeuralNetNode):
"""A leaf node that represents an input to the network"""
value: float=0.0
def op(self, x):
return self.value
def grad(self) -> List[float]:
return [1.0]
def __repr__(self):
return f'{self.__class__.__name__}({self.value})'
@dataclass
class Parameter(NeuralNetNode):
"""A leaf node that represents a parameter to the network"""
value: float=field(default_factory=lambda: random.uniform(-1, 1))
learning_rate: float=0.01
def op(self, x):
return self.value
def grad(self):
return [1.0]
def on_backward(self, derivative_from_parent: float):
self.value -= derivative_from_parent * self.learning_rate
def __repr__(self):
return f'{self.__class__.__name__}({self.value})'
Operations are barely more complicated, though not an excessive amount of so — we just have to calculate their gradients properly. Below are implementations of some useful operations:
import math@dataclass
class Operation(NeuralNetNode):
"""A node that performs an operation on its inputs"""
pass
@dataclass
class Add(Operation):
"""A node that adds its inputs"""
def op(self, x):
return sum(x)
def grad(self):
return [1.0] * len(self.children)
@dataclass
class Multiply(Operation):
"""A node that multiplies its inputs"""
def op(self, x):
return math.prod(x)
def grad(self):
grads = []
for i in range(len(self.children)):
cur_grad = 1
for j in range(len(self.children)):
if i == j:
proceed
cur_grad *= self.children[j].forward()
grads.append(cur_grad)
return grads
@dataclass
class ReLU(Operation):
"""
A node that applies the ReLU function to its input.
Note that this could only have one child.
"""
def op(self, x):
return max(0, x[0])
def grad(self):
return [1.0 if self.children[0].forward() > 0 else 0.0]
@dataclass
class Sigmoid(Operation):
"""
A node that applies the sigmoid function to its input.
Note that this could only have one child.
"""
def op(self, x):
return 1 / (1 + math.exp(-x[0]))
def grad(self):
return [self.forward() * (1 - self.forward())]
The operation superclass here will not be useful yet, though we are going to need it to more easily find our model’s inputs later.
Notice how often the gradients of the functions require the values from their children, hence we require calling the kid’s forward() method. We’ll touch upon this more in just a little bit.
Defining a neural network in our framework is a bit verbose but may be very much like constructing a tree. Here, as an illustration, is code for a straightforward linear classifier in our framework:
linear_classifier = Add([
Multiply([
Parameter(),
Input()
]),
Parameter()
])
Using Our Models
To run a prediction with our model, we now have to first populate the inputs in our tree after which call forward() on the parent. To populate the inputs though, we first need to search out them, hence we add the next method to our Operation class (we don’t add this to our NeuralNetNode class for the reason that Input type isn’t defined there yet):
def find_input_nodes(self) -> List[Input]:
"""Find all the input nodes within the subtree rooted at this node"""
input_nodes = []
for child in self.children:
if isinstance(child, Input):
input_nodes.append(child)
elif isinstance(child, Operation):
input_nodes.extend(child.find_input_nodes())
return input_nodes
We will now add the predict() method to the Operation class:
def predict(self, inputs: List[float]) -> float:
"""Evaluate the network on the given inputs"""
input_nodes = self.find_input_nodes()
assert len(input_nodes) == len(inputs)
for input_node, value in zip(input_nodes, inputs):
input_node.value = value
return self.forward()
Exercise 2: The present way we implemented predict() is somewhat inefficient since we want to traverse the tree to search out all of the inputs each time we run predict(). Write a compile() method that caches the operation’s inputs when it’s run.
Training our models is now very straightforward:
from typing import Callable, Tupledef train_model(
model: Operation,
loss_fn: Callable[[float, float], float],
loss_grad_fn: Callable[[float, float], float],
data: List[Tuple[List[float], float]],
epochs: int=1000,
print_every: int=100
):
"""Train the given model on the given data"""
for epoch in range(epochs):
total_loss = 0.0
for x, y in data:
prediction = model.predict(x)
total_loss += loss_fn(y, prediction)
model.backward(loss_grad_fn(y, prediction))
if epoch % print_every == 0:
print(f'Epoch {epoch}: loss={total_loss/len(data)}')
Here, as an illustration, is how we’d train a linear Fahrenheit to Celsius classifier using our framework:
def mse_loss(y_true: float, y_pred: float) -> float:
return (y_true - y_pred) ** 2def mse_loss_grad(y_true: float, y_pred: float) -> float:
return -2 * (y_true - y_pred)
def fahrenheit_to_celsius(x: float) -> float:
return (x - 32) * 5 / 9
def generate_f_to_c_data() -> List[List[float]]:
data = []
for _ in range(1000):
f = random.uniform(-1, 1)
data.append([[f], fahrenheit_to_celsius(f)])
return data
linear_classifier = Add([
Multiply([
Parameter(),
Input()
]),
Parameter()
])
train_model(linear_classifier, mse_loss, mse_loss_grad, generate_f_to_c_data())
After running this, we get
print(linear_classifier)
print(linear_classifier.predict([32]))>> Add(children=[Multiply(children=[Parameter(0.5555555555555556), Input(0.8930639016107234)]), Parameter(-17.777777777777782)])
>> -1.7763568394002505e-14
Which appropriately corresponds to a linear classifier with weight 0.56, bias -17.78 (which is the Fahrenheit to Celsius formula)
We will, in fact, also train rather more complex models, e.g. here is one for predicting if a degree (x, y) is above or below the road y = x:
def bce_loss(y_true: float, y_pred: float, eps: float=0.00000001) -> float:
y_pred = min(max(y_pred, eps), 1 - eps)
return -y_true * math.log(y_pred) - (1 - y_true) * math.log(1 - y_pred)def bce_loss_grad(y_true: float, y_pred: float, eps: float=0.00000001) -> float:
y_pred = min(max(y_pred, eps), 1 - eps)
return (y_pred - y_true) / (y_pred * (1 - y_pred))
def generate_binary_data():
data = []
for _ in range(1000):
x = random.uniform(-1, 1)
y = random.uniform(-1, 1)
data.append([(x, y), 1 if y > x else 0])
return data
model_binary = Sigmoid(
[
Add(
[
Multiply(
[
Parameter(),
ReLU(
[
Add(
[
Multiply(
[
Parameter(),
Input()
]
),
Multiply(
[
Parameter(),
Input()
]
),
Parameter()
]
)
]
)
]
),
Parameter()
]
)
]
)
train_model(model_binary, bce_loss, bce_loss_grad, generate_binary_data())
Then we reasonably get
print(model_binary.predict([1, 0]))
print(model_binary.predict([0, 1]))
print(model_binary.predict([0, 1000]))
print(model_binary.predict([-5, 3]))
print(model_binary.predict([0, 0]))>> 3.7310797619230176e-66
>> 0.9997781079343139
>> 0.9997781079343139
>> 0.9997781079343139
>> 0.23791579184662365
Though this has reasonable runtime, it’s somewhat slower than we’d expect. It is because we now have to call forward() and re-calculate the model inputs quite a bit in the decision to backwards(). As such, have the next exercise:
Exercise 3: Add caching to our network. That’s, on the decision to forward(), the model should return the cached value from the previous call to forward() if and provided that the inputs haven’t modified for the reason that last call. Be sure that you run forward() again if the inputs have modified.
And that’s about it! We now have a working neural network framework wherein we will train just numerous interesting models (though not networks with nodes that feed into multiple other nodes. This isn’t too difficult so as to add — see the note within the discussion of the chain rule), though granted it’s a bit verbose. When you’d wish to make it higher, try a few of the following:
Exercise 4: When you concentrate on it, more “complex” nodes in our network (e.g. Linear layers) are really just “macros” in a way — that’s, if we had a neural net tree that looked, say, as follows:
what you’re really doing is that this:
In other words, Linear(inp) is absolutely only a macro for a tree containing |inp| + 1 parameters, the primary of that are weights in multiplication and the last of which is a bias. Each time we see Linear(inp), we will thus substitute it for an equivalent tree composed only of elementary operations.
For this exercise, your job is thus to implement the Macro class. The category ought to be an Operation that recursively replaces itself with elementary operations
Note: this step might be done at any time when, though it’s likely easiest so as to add a compile() method to the Operation class that you may have to call before training (or add it to your existing method from Exercise 2). We will, in fact, also implement more complex nodes in other (perhaps more efficient) ways, but this continues to be a very good exercise.
Exercise 5: Though we don’t really ever need internal nodes to supply anything aside from one number as their output, it is typically nice for the foundation of our tree (that’s, our output layer) to supply something else (e.g. an inventory of numbers within the case of a Softmax). Implement the Output class and permit it to supply a Listof[float] as an alternative of only a float. As a bonus, try implementing the SoftMax output.
Note: there are just a few ways of doing this. You’ll be able to make Output extend Operation, after which modify the NeuralNetNode class’ op() method to return a List[float] as an alternative of only a float. Alternatively, you can create a brand new Node superclass that each Output and Operation extend. This is probably going easier.
Note further that although these outputs can produce lists, they are going to still only get one derivative back from the loss function — the loss function will just occur to take an inventory of floats as an alternative of a float (e.g. the Categorical Cross Entropy loss)
Exercise 6: Remember how earlier within the article we said that neural nets are only mathematical functions comprised of elementary operations? Add the funcify() method to the NeuralNetNode class that turns it into such a function written in human-readable notation (add parentheses as you please). For instance, the neural net Add([Parameter(0.1), Parameter(0.2)]) should collapse to “0.1 + 0.2” (or “(0.1 + 0.2)”).
Note: For this to work, inputs ought to be named. When you did exercise 2, name your inputs within the compile() function. If not, you’ll need to work out a method to name your inputs — writing a compile() function continues to be likely the simplest way.
Exercise 7: Modify our framework to permit nodes to have multiple parents. I’ll solve this partially 2.
That’s all for now! When you’d like to examine out the code, you may take a look at this google colab that has all the pieces (apart from solutions to each exercise but #6, though I’ll add those partially 2).
Contact me at mchak@calpoly.edu for any inquiries.
Unless otherwise specified, all images are by the creator.