<jupyter_start><jupyter_text>The forward and backward passes [Jump_to lesson 8 video](https://course.fast.ai/videos/?lesson=8&t=4960)<jupyter_code>%load_ext autoreload
%autoreload 2

%load_ext nb_black
# export
from exp.nb_01 import *


def get_data():
    """Download MNIST dataset and return training and validation tensors."""
    path = datasets.download_data(MNIST_URL, ext=".gz")

    with gzip.open(path, "rb") as f:
        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding="latin-1")

    return map(tensor, (x_train, y_train, x_valid, y_valid))


def normalize(x, mean, std):
    """Normalize x by first centering it and then scale it."""
    return (x - mean) / std
# Get MNIST data
x_train, y_train, x_valid, y_valid = get_data()
x_train.shape, x_valid.shape
train_mean, train_std = x_train.mean(), x_train.std()
train_mean, train_std
x_train = normalize(x_train, train_mean, train_std)
# NB: Use training, not validation mean for validation set
x_valid = normalize(x_valid, train_mean, train_std)
train_mean, train_std = x_train.mean(), x_train.std()
train_mean, train_std<jupyter_output><empty_output><jupyter_text>Now the images have a mean of 0 and a standard deviation of 1.<jupyter_code># export
def test_near_zero(a, tol=1e-3):
    assert a.abs() < tol, f"Near zero: {a}"
test_near_zero(x_train.mean())
test_near_zero(1 - x_train.std())
n, m = x_train.shape
c = y_train.max() + 1
n, m, c<jupyter_output><empty_output><jupyter_text>We have 10 classes.  Foundations version  Basic architecture<jupyter_code># num hidden
nh = 50<jupyter_output><empty_output><jupyter_text>Dividing the initialization by the square root of the number of hidden units in each layer will result in weight matrix that has mean of ~ 0 and std ~ 1 / sqrt( hidden units).<jupyter_code># simplified kaiming init / he init
w1 = torch.randn(m, nh) / math.sqrt(m)
b1 = torch.zeros(nh)
w2 = torch.randn(nh, 1) / math.sqrt(nh)
b2 = torch.zeros(1)
test_near_zero(w1.mean())
test_near_zero(w1.std() - 1 / math.sqrt(m))
# This should be ~ (0,1) (mean,std)...
x_valid.mean(), x_valid.std()
def lin(x, w, b):
    return x @ w + b
t = lin(x_valid, w1, b1)
# ...so should this, because we used kaiming init, which is designed to do this
t.mean(), t.std()
def relu(x):
    return x.clamp_min(0.0)
t = relu(lin(x_valid, w1, b1))
# ...actually it really should be this!
t.mean(), t.std()<jupyter_output><empty_output><jupyter_text>From pytorch docs: `a: the negative slope of the rectifier used after this layer (0 for ReLU by default)`$$\text{std} = \sqrt{\frac{2}{(1 + a^2) \times \text{fan_in}}}$$This was introduced in the paper that described the Imagenet-winning approach from *He et al*: [Delving Deep into Rectifiers](https://arxiv.org/abs/1502.01852), which was also the first paper that claimed "super-human performance" on Imagenet (and, most importantly, it introduced resnets!)<jupyter_code># kaiming init / he init for relu
w1 = torch.randn(m, nh) * math.sqrt(2 / m)
w1.mean(), w1.std()
t = relu(lin(x_valid, w1, b1))
t.mean(), t.std()
# export
from torch.nn import init<jupyter_output><empty_output><jupyter_text>- If we use `mode=fan_in`, this means we preserve the variance of the forward pass by dividing by sqrt(1 / m); where `m` is the number of hidden units in the input layer.- If we use `mode=fan_out`, this means we preserve the variance of the backward pass by dividing by sqrt(1 / nh); where `nh` is the number of hidden units in the output layer.<jupyter_code>w1 = torch.zeros(m, nh)
init.kaiming_normal_(w1, mode="fan_out")
t = relu(lin(x_valid, w1, b1))<jupyter_output><empty_output><jupyter_text>We are using `fan_out` instead of `fan_in` because PyTorch returns the weight nh x m since the linear layer operation is `x.matmul(w.t())`.<jupyter_code>w1.mean(), w1.std()
init.kaiming_normal_??
torch.nn.init._calculate_correct_fan??
torch.nn.Linear??
torch.nn.init._calculate_fan_in_and_fan_out??
t.mean(), t.std()
w1.shape
import torch.nn
torch.nn.Linear(m, nh).weight.shape
torch.nn.Linear.forward??
torch.nn.functional.linear??
# input.matmul(weight.t())
torch.nn.Conv2d??
torch.nn.modules.conv._ConvNd.reset_parameters??<jupyter_output><empty_output><jupyter_text>This is how PyTorch initialized the Conv layer:- `init.kaiming_uniform_(self.weight, a=math.sqrt(5))`It is not documented where the sqrt(5) comes from. The following **RELU** is based on Jeremy's experimentation and he said it works better.<jupyter_code># what if...?
def relu(x):
    return x.clamp_min(0.0) - 0.5
# kaiming init / he init for relu
w1 = torch.randn(m, nh) * math.sqrt(2.0 / m)
t1 = relu(lin(x_valid, w1, b1))
t1.mean(), t1.std()
def model(xb):
    """One hidden layer model."""
    l1 = lin(xb, w1, b1)
    l2 = relu(l1)
    l3 = lin(l2, w2, b2)
    return l3
%timeit -n 10 _ = model(x_valid)
assert model(x_valid).shape == torch.Size([x_valid.shape[0], 1])<jupyter_output><empty_output><jupyter_text>Loss function: MSE<jupyter_code>model(x_valid).shape<jupyter_output><empty_output><jupyter_text>We need `squeeze()` to get rid of that trailing (,1), in order to use `mse`. (Of course, `mse` is not a suitable loss function for multi-class classification; we'll use a better loss function soon. We'll use `mse` for now to keep things simple.)<jupyter_code># export
def mse(output, targ):
    """Compute Mean Squared Error."""
    return (output.squeeze(-1) - targ).pow(2).mean()
y_train, y_valid = y_train.float(), y_valid.float()
preds = model(x_train)
preds.shape
mse(preds, y_train)<jupyter_output><empty_output><jupyter_text>Gradients and backward pass<jupyter_code>def mse_grad(y_hat, targ):
    # grad of loss with respect to output of previous layer
    y_hat.g = 2.0 * (y_hat.squeeze() - targ).unsqueeze(-1) / y_hat.shape[0]
def relu_grad(inp, out):
    # grad of relu with respect to input activations
    inp.g = (inp > 0).float() * out.g
def lin_grad(inp, out, w, b):
    # grad of matmul with respect to input
    w.g = inp.t() @ out.g
    b.g = out.g.sum(0)
    assert w.shape == w.g.shape
    assert b.shape == b.g.shape
    inp.g = out.g @ w.t()
def forward_and_backward(inp, targ):
    # forward pass:
    l1 = inp @ w1 + b1
    l2 = relu(l1)
    out = l2 @ w2 + b2
    # we don't actually need the loss in backward!
    loss = mse(out, targ)

    # backward pass:
    mse_grad(out, targ)
    lin_grad(l2, out, w2, b2)
    relu_grad(l1, l2)
    lin_grad(inp, l1, w1, b1)
forward_and_backward(x_train, y_train)
# Save for testing against later
w1g = w1.g.clone()
w2g = w2.g.clone()
b1g = b1.g.clone()
b2g = b2.g.clone()
ig = x_train.g.clone()<jupyter_output><empty_output><jupyter_text>We cheat a little bit and use PyTorch autograd to check our results.<jupyter_code>xt2 = x_train.clone().requires_grad_(True)
w12 = w1.clone().requires_grad_(True)
w22 = w2.clone().requires_grad_(True)
b12 = b1.clone().requires_grad_(True)
b22 = b2.clone().requires_grad_(True)
def forward(inp, targ):
    # forward pass:
    l1 = inp @ w12 + b12
    l2 = relu(l1)
    out = l2 @ w22 + b22
    # we don't actually need the loss in backward!
    return mse(out, targ)
loss = forward(xt2, y_train)
loss.backward()
test_near(w22.grad, w2g)
test_near(b22.grad, b2g)
test_near(w12.grad, w1g)
test_near(b12.grad, b1g)
test_near(xt2.grad, ig)<jupyter_output><empty_output><jupyter_text>Now we are sure that our implementation of both forward and backward passes are correct after comparing it with PyTorch.  Refactor model  Layers as classes<jupyter_code>class Relu:
    def __call__(self, inp):
        self.inp = inp
        self.out = inp.clamp_min(0.0) - 0.5
        return self.out

    def backward(self):
        self.inp.g = (self.inp > 0).float() * self.out.g
class Lin:
    def __init__(self, w, b):
        self.w, self.b = w, b

    def __call__(self, inp):
        self.inp = inp
        self.out = inp @ self.w + self.b
        return self.out

    def backward(self):
        self.w.g = self.inp.t() @ self.out.g
        self.b.g = self.out.g.sum(0)
        self.inp.g = self.out.g @ self.w.t()
class Mse:
    def __call__(self, inp, targ):
        self.inp = inp
        self.targ = targ
        self.out = (self.inp.squeeze() - self.targ).pow(2).mean()
        return self.out

    def backward(self):
        self.inp.g = (
            2.0 * (self.inp.squeeze(-1) - self.targ).unsqueeze(-1) / self.targ.shape[0]
        )
class Model:
    def __init__(self, w1, b1, w2, b2):
        self.layers = [Lin(w1, b1), Relu(), Lin(w2, b2)]
        self.loss = Mse()

    def __call__(self, x, targ):
        for l in self.layers:
            x = l(x)
        return self.loss(x, targ)

    def backward(self):
        self.loss.backward()
        for l in reversed(self.layers):
            l.backward()
w1.g, b1.g, w2.g, b2.g = [None] * 4
model = Model(w1, b1, w2, b2)
%time loss = model(x_train, y_train)
%time model.backward()
test_near(w2g, w2.g)
test_near(b2g, b2.g)
test_near(w1g, w1.g)
test_near(b1g, b1.g)
test_near(ig, x_train.g)<jupyter_output><empty_output><jupyter_text>Module.forward()<jupyter_code>class Module:
    def __call__(self, *args):
        self.args = args
        self.out = self.forward(*args)
        return self.out

    def forward(self):
        raise Exception("not implemented")

    def backward(self):
        self.bwd(self.out, *self.args)
class Relu(Module):
    def forward(self, inp):
        return inp.clamp_min(0.0) - 0.5

    def bwd(self, out, inp):
        inp.g = (inp > 0).float() * out.g
class Lin(Module):
    def __init__(self, w, b):
        self.w, self.b = w, b

    def forward(self, inp):
        return inp @ self.w + self.b

    def bwd(self, out, inp):
        self.w.g = inp.t() @ out.g
        self.b.g = out.g.sum(0)
        inp.g = out.g @ self.w.t()
class Mse(Module):
    def forward(self, inp, targ):
        return (inp.squeeze() - targ).pow(2).mean()

    def bwd(self, out, inp, targ):
        inp.g = 2 * (inp.squeeze() - targ).unsqueeze(1) / targ.shape[0]
class Model:
    def __init__(self):
        self.layers = [Lin(w1, b1), Relu(), Lin(w2, b2)]
        self.loss = Mse()

    def __call__(self, x, targ):
        for l in self.layers:
            x = l(x)
        return self.loss(x, targ)

    def backward(self):
        self.loss.backward()
        for l in reversed(self.layers):
            l.backward()
w1.g, b1.g, w2.g, b2.g = [None] * 4
model = Model()
%time loss = model(x_train, y_train)
%time model.backward()
test_near(w2g, w2.g)
test_near(b2g, b2.g)
test_near(w1g, w1.g)
test_near(b1g, b1.g)
test_near(ig, x_train.g)<jupyter_output><empty_output><jupyter_text>nn.Linear and nn.Module<jupyter_code># export
from torch import nn
class Model(nn.Module):
    def __init__(self, n_in, nh, n_out):
        super().__init__()
        self.layers = [nn.Linear(n_in, nh), nn.ReLU(), nn.Linear(nh, n_out)]
        self.loss = mse

    def __call__(self, x, targ):
        for l in self.layers:
            x = l(x)
        return self.loss(x.squeeze(), targ)
model = Model(m, nh, 1)
%time loss = model(x_train, y_train)
%time loss.backward()<jupyter_output>CPU times: user 113 ms, sys: 3.16 ms, total: 116 ms
Wall time: 28.9 ms<jupyter_text>Export<jupyter_code>!python ../src/notebook2script.py 02_Fully-Connected.ipynb<jupyter_output>Converted 02_Fully-Connected.ipynb to exp/nb_02.py
{'fraction_non_alphanumeric': 0.14297117516629712, 'fraction_numerical': 0.022350332594235033, 'mean_word_length': 3.5688816855753647, 'pattern_counts': {'":': 0, '<': 64, '<?xml version=': 0, '>': 66, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '16', 'dir': 'jupyter-structured-clean-dedup', 'id': '283769', 'license': 'Apache-2.0', 'n_tokens_mistral': 4545, 'n_tokens_neox': 4384, 'n_words': 1132, 'path': 'dl2/notebooks/02_Fully-Connected.ipynb', 'repo_name': 'ImadDabbura/fastai-courses', 'size': '11275'}