# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + id="GJ-wL42Lxp7g"
# %load_ext autoreload
# %autoreload 2

# %matplotlib inline

# + id="AO3i1A7iL--k"

#from google.colab import drive
#drive.mount('/content/gdrive')

# + colab={"base_uri": "https://localhost:8080/"} id="WdH_c8nnqn-c" outputId="0dba1de3-0b50-4c27-c757-29d573843927"
#Q1. Execute the following statement. What is displayed? What does it mean?
# !pwd

# + id="jhAroxykqtRF"
#Q2: Execute the following statement. What happens?  Examine the left column of your colab page to see what happens.
# !git clone https://github.com/fastai/course-v3.git

# + id="FOZb4AeKq6vQ"
#Q3. Execute the following statement. What are displayed by the execution of sys.path? What do it mean?
import sys
sys.path.append('/content/course-v3/nbs/dl2')
sys.path

# + id="I86djNvcxp70"
#export
#Q4: "from exp.nb_01 import *" imports classes, functions, etc from module nb_01 in package exp.
# Q4a.   What do you know that exp is a package?
# Q4b.   How is the Python system able to get access to package exp? That is, how can it search for it?
from exp.nb_02 import *
import torch.nn.functional as F

# + [markdown] id="iywbFmpbxp73"
# ## Initial setup

# + [markdown] id="G5DBpkZWxp75"
# ### Data

# + [markdown] id="Q_7vJLX_xp75"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=1786)

# + id="xg8JmuuDLwmt"
from IPython.display import Image
from six.moves import urllib

opener = urllib.request.build_opener()
opener.addheaders = [('User-agent', 'Mozilla/5.0')]
urllib.request.install_opener(opener)

def get_data():
    import os
    import torchvision.datasets as datasets
    root = '../data'
    if not os.path.exists(root):
        os.mkdir(root)

    train_set = datasets.MNIST(root=root, train=True, download=True)
    test_set = datasets.MNIST(root=root, train=False, download=True) #load validation set
    x_train, x_valid = train_set.data.split([50000, 10000])
    y_train, y_valid = train_set.targets.split([50000, 10000])
    return (x_train.view(50000, -1) / 256.0), y_train.float(), (x_valid.view(10000, -1))/ 256.0, y_valid.float()



# + id="TRKV-q7Kxp76"
mpl.rcParams['image.cmap'] = 'gray'

# + id="i6um-ZEMxp76"
#Q5: when you execute the following statement, where is the downloaded data stored? Examine the left column of your colab page.
x_train,y_train,x_valid,y_valid = get_data()

# + colab={"base_uri": "https://localhost:8080/"} id="dwBgUYjX1Gbc" outputId="7aa49d60-0624-4d7a-d29f-b57fc9aae53b"
y_train.shape

# + id="oHcenWs_rZ_I"
#Q6: Execute the following statement. What does the number displayed mean?
len(x_train)

# + id="mxSwa7g5rcer"


# + id="DESekbEYxp77"
#Q7: Display the values of n,m, c, and nh. What are they? For what are they used in the following code? 
n,m = x_train.shape
c = y_train.max()+1
nh = 50


# + id="_cCzC9iZxp77"
#Q8: The following defines Model class, which will be used to create a neural net.
#Q8a: What is nn.Module?, What is nn.Linear? What is nn.ReLU? 
#Q8b. What is the difference between  nn.Linear itself and nn.Linear(n_in,nh) ? 

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)]
        
    def __call__(self, x):
        for l in self.layers: x = l(x)
        return x


# + id="LJx4LIA1xp78"
#Q9a. When Model(m,nh, 10) is executed, which method in Model class is executed?
model = Model(m, nh, 10)

# + id="gPFt346vxp78"
#Q10a. Is model a function or an object? If it is an object, how is it used as if it is a function?
# When model(x_train) is executed, which method in Model is executed. What does this method create? 
pred = model(x_train)

# + id="gXbLts9ZRMgd"
pred

# + id="FsWmZbCWRQd_"
#Q11. What does pred[0] contain?
pred[0]


# + id="-Ejrm01LRSxj"
#Q12. Execute the following statement. Is the resulting value near to 1? If not, what does it imply?
pred[0].sum()


# + [markdown] id="qhbGAsGWxp79"
# ### Cross entropy loss

# + [markdown] id="fFGQT0zpxp79"
# First, we will need to compute the softmax of our activations. This is defined by:
#
# $$\hbox{softmax(x)}_{i} = \frac{e^{x_{i}}}{e^{x_{0}} + e^{x_{1}} + \cdots + e^{x_{n-1}}}$$
#
# or more concisely:
#
# $$\hbox{softmax(x)}_{i} = \frac{e^{x_{i}}}{\sum_{0 \leq j \leq n-1} e^{x_{j}}}$$ 
#
# In practice, we will need the log of the softmax when we calculate the loss.

# + id="Q3ZRxbanxp7_"
def log_softmax(x): return (x.exp()/(x.exp().sum(-1,keepdim=True))).log()


# + id="QDeeKw-Sxp8A"
log_sm_pred = log_softmax(pred)

# + colab={"base_uri": "https://localhost:8080/"} id="PJmGXhPfRYba" outputId="5678501b-3553-42cf-f595-3d7c32de04e4"
log_sm_pred[0]

# + id="JRmNSz1vRb0N"
log_sm_pred[0].sum()

# + [markdown] id="DfNuX2v1xp8B"
# ## Log likelihood interpretation: https://datascience.stackexchange.com/questions/9302/the-cross-entropy-error-function-in-neural-networks
#
# One way to interpret cross-entropy is to see it as a (minus) log-likelihood for the data y′i, under a model yi.
#
# Namely, suppose that you have some fixed model (a.k.a. "hypothesis"), which predicts for n classes {1,2,…,n} their hypothetical occurrence probabilities y1,y2,…,yn. Suppose that you now observe (in reality) k1 instances of class 1, k2 instances of class 2, kn instances of class n, etc. According to your model the likelihood of this happening is:
#
# $P[data|model]:=y_{1}^{k_{1}}...y_{n}^{k_{n}}$ .
#
# Taking the logarithm and changing the sign:
# −logP[data|model]= 
# $−k_{1} * logy_{1}−k_{2} * logy_{2}−⋯−k_{n}logy_{n} =−\sum_{i} k_{i}*logy_{i}$
#
# If you now divide the right-hand sum by the number of observations N=k1+k2+⋯+kn, and denote the empirical probabilities as y′i=ki/N, you'll get the cross-entropy:
# −1/N * logP[data|model]=−1/N
# $\sum_{i}k_{i}logy_{i}=−\sum_{i}y′_{i}logy_{i} = H(y′,y)$
#
# ## The Differnce between probabilities interpretation: https://datascience.stackexchange.com/questions/20296/cross-entropy-loss-explanation
#
# $H(p,\hat{p})=−\sum_{s \in B} p(x^{(s)}) \cdot log(\hat{p}(x^{(s)}))$
#
#
# In general, $ p(x^{(s)}) = [ p_{1} (x^{(s)}), ..., p_{n}(x^{(s)})]$ is a probability distribution over a set of categories. 
# But since our $ p(x^{(s)})$ are 1-hot encoded, that is, in the form of  $ p(x^{(s)}) =[0,0,..,0,1,0..]$, where the probability of only one category is one and those of the other categories are all zero, this can be rewritten as
#
#  $-\sum_{s \in B} [ 0*\log(\hat{p}_{1} (x^{(s)}))
#  + 1*\log(\hat{p}_{i} (x^{(s)}) ) +..+0*\log(\hat{p}_{n} (x^{(s)}) ) ] \\
#   = -\sum_{s \in B} \log(\hat{p}_{i(s)} ) (x^{(s)}) )$ 
#
#   Here $i(s)$ is the index of the one-hot  probability distribution $p(x^{(s)})$ where the probability is one. 
#  
#  

# + [markdown] id="KBtZhgePxp8D"
# This can be done using numpy-style [integer array indexing](https://docs.scipy.org/doc/numpy-1.13.0/reference/arrays.indexing.html#integer-array-indexing). Note that PyTorch supports all the tricks in the advanced indexing methods discussed in that link.

# + id="3R_y0cP_xp8D" colab={"base_uri": "https://localhost:8080/"} outputId="0df5c3d0-18c7-4a14-870a-226f7ad09a24"
y_train[:3]

# + colab={"base_uri": "https://localhost:8080/"} id="VehY657jB7IB" outputId="9937fcbc-4ddb-43ed-a6fd-852936cc8d63"
log_sm_pred[ [0,1,2]]

# + id="uJM_9p0rxp8H" colab={"base_uri": "https://localhost:8080/"} outputId="929485f1-54f3-4291-cc50-cdb67c2e1a89"
log_sm_pred[[0,1,2], [5,0,4]]

# + id="oSK1T-mgxp8I"
y_train.shape[0]

# + colab={"base_uri": "https://localhost:8080/"} id="-ogBL3jZxPPN" outputId="01119f33-9e56-4fcd-e7a8-5cdad566afe1"
y_train

# + [markdown] id="h1bTmfz5xp8J"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=2081)

# + colab={"base_uri": "https://localhost:8080/"} id="pPJ9QjlLQQc_" outputId="7323044b-38e6-4359-a4c5-183343ec16f8"
log_sm_pred[0]

# + colab={"base_uri": "https://localhost:8080/"} id="xoc2YwU5wATE" outputId="db5c56ac-c45e-4815-9908-ded483996b71"
range(y_train.long().shape[0])

# + colab={"base_uri": "https://localhost:8080/"} id="8GMBhmq9w4QG" outputId="d18ff032-6585-4755-9142-e6ff64d13682"
log_sm_pred[range(y_train.long().shape[0]), y_train[0:10].long() ].shape


# + id="ZI0cXSLRxp8K"
def nll(input, target):
   
   return -input[range(target.shape[0]), target].mean()


# + colab={"base_uri": "https://localhost:8080/"} id="vzmLXXWiPTSF" outputId="a8bf7fd5-d0d3-4669-8bad-3ddfa532dac8"
y_train.long().shape[0]

# + colab={"base_uri": "https://localhost:8080/"} id="_qRbZvGePbHV" outputId="5d8c06da-0da9-4699-f69a-c65b3b91a183"
x_train.shape

# + id="Bs5UU4fyxp8K"
loss = nll(log_sm_pred, y_train.long())

# + colab={"base_uri": "https://localhost:8080/"} id="rrPFVa91P0EO" outputId="522c4008-bfcb-486c-9f69-9aa19ee2bcf3"
y_train

# + colab={"base_uri": "https://localhost:8080/"} id="6uzxkFQaPv-W" outputId="f229a792-a57f-45a9-d73a-98bf5820ddf5"
y_train.long()

# + id="zkSkgQYQxp8L" colab={"base_uri": "https://localhost:8080/"} outputId="452d86ad-a756-4d0d-aa4f-c84ac6c5e235"
loss


# + [markdown] id="rHyuP5t7xp8M"
# Note that the formula 
#
# $$\log \left ( \frac{a}{b} \right ) = \log(a) - \log(b)$$ 
#
# gives a simplification when we compute the log softmax, which was previously defined as `(x.exp()/(x.exp().sum(-1,keepdim=True))).log()`

# + id="veHLG6DXxp8M"
def log_softmax(x): return x - x.exp().sum(-1,keepdim=True).log()


# + id="aelGhkV5xp8N"
test_near(nll(log_softmax(pred), y_train.long()), loss)


# + [markdown] id="lVPZMaTcxp8O"
# Then, there is a way to compute the log of the sum of exponentials in a more stable way, called the [LogSumExp trick](https://en.wikipedia.org/wiki/LogSumExp). The idea is to use the following formula:
#
# $$\log \left ( \sum_{j=1}^{n} e^{x_{j}} \right ) = \log \left ( e^{a} \sum_{j=1}^{n} e^{x_{j}-a} \right ) = a + \log \left ( \sum_{j=1}^{n} e^{x_{j}-a} \right )$$
#
# where a is the maximum of the $x_{j}$.

# + id="No5iHYvCxp8P"
def logsumexp(x):
    m = x.max(-1)[0]
    return m + (x-m[:,None]).exp().sum(-1).log()


# + [markdown] id="UpzazYXDxp8Q"
# This way, we will avoid an overflow when taking the exponential of a big activation. In PyTorch, this is already implemented for us. 

# + id="dQ_VQU0Sxp8Q"
test_near(logsumexp(pred), pred.logsumexp(-1))


# + [markdown] id="wkGCt3oDxp8R"
# So we can use it for our `log_softmax` function.

# + id="q7nsaEFtxp8R"
def log_softmax(x): return x - x.logsumexp(-1,keepdim=True)


# + id="Wt1PreQAxp8S"
test_near(nll(log_softmax(pred), y_train.long()), loss)

# + [markdown] id="FcdHBuPOxp8S"
# Then use PyTorch's implementation.

# + id="U609ws15xp8T"
test_near(F.nll_loss(F.log_softmax(pred, -1), y_train.long()), loss)

# + [markdown] id="yOo22f7Exp8U"
# In PyTorch, `F.log_softmax` and `F.nll_loss` are combined in one optimized function, `F.cross_entropy`.

# + id="lU7uCHBwxp8U"
test_near(F.cross_entropy(pred, y_train.long()), loss)

# + [markdown] id="BdFUzc94xp8V"
# ## Basic training loop

# + [markdown] id="495DuZ5Xxp8V"
# Basically the training loop repeats over the following steps:
# - get the output of the model on a batch of inputs
# - compare the output to the labels we have and compute a loss
# - calculate the gradients of the loss with respect to every parameter of the model
# - update said parameters with those gradients to make them a little bit better

# + [markdown] id="6e6Y1AKyxp8W"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=2542)

# + id="fDHaEy3oxp8X"
loss_func = F.cross_entropy # This criterion combines log_softmax and nll_loss in a single function.


# + id="wcuWlwDXxp8Y"
#export
def accuracy(out, yb): return (torch.argmax(out, dim=1)==yb).float().mean()


# + id="ozBbmfr7xp8Z" colab={"base_uri": "https://localhost:8080/"} outputId="ecc22746-beca-402f-9126-fc1f10c21c9e"
bs=64                  # batch size

xb = x_train[0:bs]     # a mini-batch from x
preds = model(xb)      # predictions
preds[0], preds.shape

# + id="Ue2zaDOLxp8a" colab={"base_uri": "https://localhost:8080/"} outputId="7c2ce71e-5f8f-42d4-a0ad-2dd8b1fcd073"
yb = y_train[0:bs]
loss_func(preds, yb.long())

# + id="vYSsK_Rlxp8b" colab={"base_uri": "https://localhost:8080/"} outputId="856100ed-dc6e-4991-ebfb-d78b3618eb6a"
accuracy(preds, yb)

# + id="dxLMr6u-xp8c"
lr = 0.5   # learning rate
epochs = 1 # how many epochs to train for

# + [markdown] id="9-hKTgAOFvu0"
# # The mechanism for training the network: (1) computing the graidents of Loss with respect to tensors at each layer, (2) backpropgation: applying the chain rule to compute the gradient vector of the parameters, (3) updating the parameters.
#
# https://github.com/pytorch/pytorch/blob/35bd2b3c8b64d594d85fc740e94c30aa67892a34/torch/tensor.py
#
# https://github.com/pytorch/pytorch/blob/35bd2b3c8b64d594d85fc740e94c30aa67892a34/torch/tensor.py
#
# https://stackoverflow.com/questions/57248777/backward-function-in-pytorch
# : Pytorch does not support this non-scalar function derivatives. Instead, pytorch assumes out is only an intermediate tensor and somewhere "upstream" there is a scalar loss function, that through chain rule provides d loss/ d out[i,j]. This "upstream" gradient is of size 2-by-3 and this is actually the argument you provide backward in this case: out.backward(g) where g_ij = d loss/ d out_ij.
#
# The gradients are then calculated by chain rule d loss / d a[i,j] = (d loss/d out[i,j]) * (d out[i,j] / d a[i,j]); In fact,  So it's not just (d loss/d out[i,j]) * (d out[i,j] / d a[i,j]), but in fact sum_{k,l} (d loss/d out[k,l]) * (d out[k,l] / d a[i,j]).
#
# # The "gradient accumulation"
#
# https://discuss.pytorch.org/t/why-do-we-need-to-set-the-gradients-manually-to-zero-in-pytorch/4903/9
#
# In the following example, y.backward() is called 5 times, so the final value of x.grad will be 5*cos(0)=5.
#
# import torch
# from torch.autograd import Variable
#
# x = Variable(torch.Tensor([[0]]), requires_grad=True)
#
# for t in range(5):
#     y = x.sin() 
#     y.backward()
#     
# print(x.grad) # shows 5
# Calling x.grad.data.zero_() before y.backward() can make sure x.grad is exactly the same as current y’(x), not a sum of y’(x) in all previous iterations.
#
# x = Variable(torch.Tensor([[0]]), requires_grad=True) 
#
# for t in range(5):
#     if x.grad is not None:
#         x.grad.data.zero_()
#     y = x.sin() 
#     y.backward()
#
# print(x.grad) # shows 1
# I also got confused by this “zeroing gradient” when first learning pytorch. The doc of torch.autograd.backward does mention that
#
# This function accumulates gradients in the leaves - you might need to zero them before calling it.
#
# But this is quite hard to find and pretty confusing for (say) tensorflow users.
#
# Official tutorials like 60 Minute Blitz 45 or PyTorch with Examples 69 both say nothing about why one needs to call grad.data.zero_() during training. I think it would be useful to explain this a little more in beginner-level tutorials. RNN is a good example for why accumulating gradient (instead of refreshing) is useful, but I guess new users wouldn’t even know that backward() is accumulating gradient
#
#
# # A more rigorous reference to autograd:
#
# https://blog.paperspace.com/pytorch-101-understanding-graphs-and-automatic-differentiation/
#
# torch.no_grad():
#
# When we are computing gradients, we need to cache input values, and intermediate features as they maybe required to compute the gradient later.
#
# The gradient of  
# b
# =
# w
# 1
# ∗
# a
#  w.r.t it's inputs 
# w
# 1
#  and 
# a
#  is 
# a
#  and 
# w
# 1
#  respectively. We need to store these values for gradient computation during the backward pass. This affects the memory footprint of the network.
#
# While, we are performing inference, we don't compute gradients, and thus, don't need to store these values. Infact, no graph needs to be create during inference as it will lead to useless consumption of memory.
#
# PyTorch offers a context manager, called torch.no_grad for this purpose.
#
#
# with torch.no_grad:
#
# 	inference code goes here 
#
# No graph is defined for operations executed under this context manager.
#
#
#

# + id="Q0KH3m-Uxp8c" colab={"base_uri": "https://localhost:8080/"} outputId="10ebc4d7-40b1-48f7-b7c3-3a15095924e5"
for epoch in range(epochs):
    for i in range((n-1)//bs + 1): # for each batch in the current epoch
#         set_trace()
        start_i = i*bs
        end_i = start_i+bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]

        ybHat = model(xb)
        h = ybHat.register_hook( lambda grad: print (grad) )

        loss = loss_func( ybHat, yb.long())

        # r"""Computes the gradient of loss w.r.t. graph LEAVES.
        #  This function accumulates gradients in the leaves - you might need to zero
        # ``.grad`` attributes or set them to ``None`` before calling it.
        #  (This makes it easier for rnn, because each module will be back propogated through several times.)
       
       # Attribute grad of tensor ybHat (ybHat.grad = dloss/dybHat) is ``None`` by default and becomes a Tensor the first time a call to
        #func `backward` computes gradients for ybHat.
        #The attribute will then contain the gradients computed and future calls to
        #:func:`backward` will accumulate (add) gradients into it.

        #Since the backward() function accumulates gradients, and you don’t want to mix up gradients between minibatches, 
        # you have to zero them out at the start of a new minibatch.
        # By the way, the best practice is to use the optimzer.zero_grad()


        ybHat.grad
        loss.backward() # When you call loss.backward(), all it does is compute gradient of loss 
                        # w.r.t all the parameters in loss that have requires_grad = True and store them in parameter.grad attribute for every parameter

        ybHat.grad                
        #optimizer.step() updates all the parameters based on parameter.grad 
        with torch.no_grad():
            for l in model.layers:
                if hasattr(l, 'weight'):
                    l.weight -= l.weight.grad * lr
                    l.bias   -= l.bias.grad   * lr

                    #we don't care about gradients from the previous batch. Not zeroing grads would lead to gradient accumulation across batches
                    l.weight.grad.zero_()
                    l.bias  .grad.zero_()

# + id="I0fC9kbvxp8g" colab={"base_uri": "https://localhost:8080/"} outputId="c8b375c1-825a-4a92-c027-49b59a68a658"
loss_func(model(xb), yb.long()), accuracy(model(xb), yb)


# + [markdown] id="n70culC-xp8h"
# ## Using parameters and optim

# + [markdown] id="Ztmln-E0xp8h"
# ### Parameters

# + [markdown] id="3JMWX2-fxp8i"
# Use `nn.Module.__setattr__` and move relu to functional:

# + [markdown] id="2qu-R0Ljxp8j"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=2818)

# + id="lw5k0KJRxp8j"
class Model(nn.Module):
    def __init__(self, n_in, nh, n_out):
        super().__init__()
        self.l1 = nn.Linear(n_in,nh)
        self.l2 = nn.Linear(nh,n_out)
        
    def __call__(self, x): return self.l2(F.relu(self.l1(x)))


# + id="lq8zwnBlxp8k"
model = Model(m, nh, 10)

# + id="yFsAG7KCxp8k" colab={"base_uri": "https://localhost:8080/"} outputId="9d765f98-3b16-4a38-98b5-6bcc457da4f2"
for name,l in model.named_children(): print(f"{name}: {l}")

# + id="QQAcSLAVxp8l" colab={"base_uri": "https://localhost:8080/"} outputId="c1301582-6aa8-4bb1-9e3b-c9ad0f685747"
model

# + id="O77Ds-LRxp8m" colab={"base_uri": "https://localhost:8080/"} outputId="ae80b6f2-ead8-43f0-a448-3888e712dd07"
model.l1


# + id="SUDx09pkxp8m"
def fit():
    for epoch in range(epochs):
        for i in range((n-1)//bs + 1):
            start_i = i*bs
            end_i = start_i+bs
            xb = x_train[start_i:end_i]
            yb = y_train[start_i:end_i]
            loss = loss_func(model(xb), yb.long())

            loss.backward()
            with torch.no_grad():
                for p in model.parameters(): p -= p.grad * lr
                model.zero_grad()


# + id="zkjGdktlxp8o" colab={"base_uri": "https://localhost:8080/"} outputId="b2db913c-a4b9-4a7d-9003-28aadf55efff"
fit()
loss_func(model(xb), yb.long()), accuracy(model(xb), yb)


# + [markdown] id="-wQ0CkLmxp8p"
# Behind the scenes, PyTorch overrides the `__setattr__` function in `nn.Module` so that the submodules you define are properly registered as parameters of the model.

# + id="nfLTjkNJxp8p"
class DummyModule():
    def __init__(self, n_in, nh, n_out):
        self._modules = {}
        self.l1 = nn.Linear(n_in,nh)
        self.l2 = nn.Linear(nh,n_out)
        
    def __setattr__(self,k,v):
        if not k.startswith("_"): self._modules[k] = v
        super().__setattr__(k,v)
        
    def __repr__(self): return f'{self._modules}'
    
    def parameters(self):
        for l in self._modules.values():
            for p in l.parameters(): yield p


# + id="BxdtM1ilxp8q" colab={"base_uri": "https://localhost:8080/"} outputId="6a4c6d65-f870-47a5-9ffb-4878cb76a6e9"
mdl = DummyModule(m,nh,10)
mdl

# + id="asWZtQT7xp8r" colab={"base_uri": "https://localhost:8080/"} outputId="bc341473-2c04-4764-dc64-65061f646c72"
[o.shape for o in mdl.parameters()]

# + [markdown] id="PJzROcD-xp8r"
# ### Registering modules

# + [markdown] id="Gi6Bj3Egxp8s"
# We can use the original `layers` approach, but we have to register the modules.

# + [markdown] id="SCDxRdQexp8s"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=2997)

# + id="8ax355bwxp8t"
layers = [nn.Linear(m,nh), nn.ReLU(), nn.Linear(nh,10)]


# + id="0FVOwYH8xp8t"
class Model(nn.Module):
    def __init__(self, layers):
        super().__init__()
        self.layers = layers
        for i,l in enumerate(self.layers): self.add_module(f'layer_{i}', l)
        
    def __call__(self, x):
        for l in self.layers: x = l(x)
        return x


# + id="x5hq-X7Vxp8u"
model = Model(layers)

# + id="lqYP2MXlxp8u" colab={"base_uri": "https://localhost:8080/"} outputId="2acdf4d1-9aee-4e4e-bf4f-5fb1488c41b1"
model


# + [markdown] id="diwEDyzJxp8v"
# ### nn.ModuleList

# + [markdown] id="yUVBUAqqxp8v"
# `nn.ModuleList` does this for us.

# + [markdown] id="w39ws0evxp8w"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=3173)

# + id="YNrrTzq6xp8w"
class SequentialModel(nn.Module):
    def __init__(self, layers):
        super().__init__()
        self.layers = nn.ModuleList(layers)
        
    def __call__(self, x):
        for l in self.layers: x = l(x)
        return x


# + id="m59ypeCdxp8x"
model = SequentialModel(layers)

# + id="SVSNzdUmxp8x" colab={"base_uri": "https://localhost:8080/"} outputId="4be69e8d-f133-4c01-8617-44e8451cae20"
model

# + id="QheaMa-6xp8y" colab={"base_uri": "https://localhost:8080/"} outputId="b02d5cb4-cbd1-4f08-ae29-1be7c5332c91"
fit()
loss_func(model(xb), yb.long()), accuracy(model(xb), yb)

# + [markdown] id="yk8EQfB6xp8y"
# ### nn.Sequential

# + [markdown] id="twIiTxJAxp8y"
# `nn.Sequential` is a convenient class which does the same as the above:

# + [markdown] id="eWOOJnAXxp8y"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=3199)

# + id="YVIaRdsPxp8z"
model = nn.Sequential(nn.Linear(m,nh), nn.ReLU(), nn.Linear(nh,10))

# + id="V_079mLyxp8z" colab={"base_uri": "https://localhost:8080/"} outputId="3994c698-7ae4-47ee-f9cf-291c194a3ce8"
fit()
loss_func(model(xb), yb.long()), accuracy(model(xb), yb)

# + id="fX3SUZc5xp8z"
# nn.Sequential??

# + id="s1EvUMznxp81" colab={"base_uri": "https://localhost:8080/"} outputId="c5637639-9c72-4336-ff2d-686a0bcb4816"
model


# + [markdown] id="1o20CTI0xp82"
# ### optim

# + [markdown] id="_8edAASgxp82"
# Let's replace our previous manually coded optimization step:
#
# ```python
# with torch.no_grad():
#     for p in model.parameters(): p -= p.grad * lr
#     model.zero_grad()
# ```
#
# and instead use just:
#
# ```python
# opt.step()
# opt.zero_grad()
# ```

# + [markdown] id="lV7oQXBxxp83"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=3278)

# + id="rz8BYVQkxp84"
class Optimizer():
    def __init__(self, params, lr=0.5): self.params,self.lr=list(params),lr
        
    def step(self):
        with torch.no_grad():
            for p in self.params: p -= p.grad * lr

    def zero_grad(self):
        for p in self.params: p.grad.data.zero_()


# + id="Rr97GcPPxp84"
model = nn.Sequential(nn.Linear(m,nh), nn.ReLU(), nn.Linear(nh,10))

# + id="lh_YE0ogxp84"
opt = Optimizer(model.parameters())

# + id="UOZPZ1rsxp85"
for epoch in range(epochs):
    for i in range((n-1)//bs + 1):
        start_i = i*bs
        end_i = start_i+bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb.long())

        loss.backward()
        opt.step()
        opt.zero_grad()

# + id="1K8AuK0yxp85" colab={"base_uri": "https://localhost:8080/"} outputId="96477a22-437a-4d1e-cedc-81ccfdc7ad38"
loss,acc = loss_func(model(xb), yb.long()), accuracy(model(xb), yb)
loss,acc

# + [markdown] id="vtEnqOylxp86"
# PyTorch already provides this exact functionality in `optim.SGD` (it also handles stuff like momentum, which we'll look at later - except we'll be doing it in a more flexible way!)

# + id="1EbDEWXsxp87"
#export
from torch import optim


# + id="DVc-OgiExp87"
# optim.SGD.step??

# + id="wJyAMVzbxp88"
def get_model():
    model = nn.Sequential(nn.Linear(m,nh), nn.ReLU(), nn.Linear(nh,10))
    return model, optim.SGD(model.parameters(), lr=lr)


# + id="fu_n5Z2zxp89" colab={"base_uri": "https://localhost:8080/"} outputId="8de39113-a531-428e-f89e-e70459000741"
model,opt = get_model()
loss_func(model(xb), yb.long())

# + id="XT0H_e6Vxp8-"
for epoch in range(epochs):
    for i in range((n-1)//bs + 1):
        start_i = i*bs
        end_i = start_i+bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb.long())

        loss.backward()
        opt.step()
        opt.zero_grad()

# + id="Do5mQXdWxp8-" colab={"base_uri": "https://localhost:8080/"} outputId="4289d3b3-4da9-498d-b2b9-96e5e69cc7dc"
loss,acc = loss_func(model(xb), yb.long()), accuracy(model(xb), yb)
loss,acc

# + [markdown] id="UVaK2RCHxp8_"
# Randomized tests can be very useful.

# + [markdown] id="TSDu_zpXxp8_"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=3442)

# + id="3GQE3EvMxp9A"
assert acc>0.7


# + [markdown] id="1PWzGE8Rxp9A"
# ## Dataset and DataLoader

# + [markdown] id="nz8ts0aZxp9A"
# ### Dataset

# + [markdown] id="hhmhOE2xxp9B"
# It's clunky to iterate through minibatches of x and y values separately:
#
# ```python
#     xb = x_train[start_i:end_i]
#     yb = y_train[start_i:end_i]
# ```
#
# Instead, let's do these two steps together, by introducing a `Dataset` class:
#
# ```python
#     xb,yb = train_ds[i*bs : i*bs+bs]
# ```

# + [markdown] id="U-gLw07Fxp9B"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=3578)

# + id="QdCJfv5Zxp9B"
#export
class Dataset():
    def __init__(self, x, y): self.x,self.y = x,y
    def __len__(self): return len(self.x)
    def __getitem__(self, i): return self.x[i],self.y[i]


# + id="5wzhoIAExp9C"
train_ds,valid_ds = Dataset(x_train, y_train),Dataset(x_valid, y_valid)
assert len(train_ds)==len(x_train)
assert len(valid_ds)==len(x_valid)

# + id="BxsllyZ8xp9C" colab={"base_uri": "https://localhost:8080/"} outputId="5a240c6a-3e35-40d2-8956-49fe539d3d28"
xb,yb = train_ds[0:5]
assert xb.shape==(5,28*28)
assert yb.shape==(5,)
xb,yb

# + id="PNvAKEUAxp9C"
model,opt = get_model()

# + id="6y3T2ONQxp9D"
for epoch in range(epochs):
    for i in range((n-1)//bs + 1):
        xb,yb = train_ds[i*bs : i*bs+bs]
        pred = model(xb)
        loss = loss_func(pred, yb.long())

        loss.backward()
        opt.step()
        opt.zero_grad()

# + id="eiGSBU34xp9D" colab={"base_uri": "https://localhost:8080/"} outputId="16d2e73b-d668-4c50-8025-5a3ae2ea8682"
loss,acc = loss_func(model(xb), yb.long()), accuracy(model(xb), yb)
assert acc>0.7
loss,acc


# + [markdown] id="0fgmSsJdxp9D"
# ### DataLoader

# + [markdown] id="QlMsRaK_xp9D"
# Previously, our loop iterated over batches (xb, yb) like this:
#
# ```python
# for i in range((n-1)//bs + 1):
#     xb,yb = train_ds[i*bs : i*bs+bs]
#     ...
# ```
#
# Let's make our loop much cleaner, using a data loader:
#
# ```python
# for xb,yb in train_dl:
#     ...
# ```

# + [markdown] id="Ka7lGsJoxp9E"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=3674)

# + id="DcOXwOi-xp9F"
class DataLoader():
    def __init__(self, ds, bs): self.ds,self.bs = ds,bs
    def __iter__(self):
        for i in range(0, len(self.ds), self.bs): yield self.ds[i:i+self.bs]


# + id="z1FCkkcdxp9F"
train_dl = DataLoader(train_ds, bs)
valid_dl = DataLoader(valid_ds, bs)

# + id="dMXh88dvxp9G"
xb,yb = next(iter(valid_dl))
assert xb.shape==(bs,28*28)
assert yb.shape==(bs,)

# + id="mWnYMKDhxp9G" colab={"base_uri": "https://localhost:8080/", "height": 283} outputId="12b6abb1-0f8c-4dac-ba8c-0ffbcee1bd21"
plt.imshow(xb[0].view(28,28))
yb[0]

# + id="mekhrFC3xp9G"
model,opt = get_model()


# + id="jIu_Wvwkxp9H"
def fit():
    for epoch in range(epochs):
        for xb,yb in train_dl:
            pred = model(xb)
            loss = loss_func(pred, yb.long())
            loss.backward()
            opt.step()
            opt.zero_grad()


# + id="JIFF7_-xxp9H"
fit()

# + id="RH5KjHAaxp9H" colab={"base_uri": "https://localhost:8080/"} outputId="6b7502d5-445e-465e-a214-710cb6fa098d"
loss,acc = loss_func(model(xb), yb.long()), accuracy(model(xb), yb)
assert acc>0.7
loss,acc


# + [markdown] id="-9uDi1fcxp9H"
# ### Random sampling

# + [markdown] id="JRARDaX9xp9P"
# We want our training set to be in a random order, and that order should differ each iteration. But the validation set shouldn't be randomized.

# + [markdown] id="MUSSX0iNxp9P"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=3942)

# + id="mgUen1XGxp9Q"
class Sampler():
    def __init__(self, ds, bs, shuffle=False):
        self.n,self.bs,self.shuffle = len(ds),bs,shuffle
        
    def __iter__(self):
        self.idxs = torch.randperm(self.n) if self.shuffle else torch.arange(self.n)
        for i in range(0, self.n, self.bs): yield self.idxs[i:i+self.bs]


# + id="nV78WqQbxp9Q"
small_ds = Dataset(*train_ds[:10])

# + id="Y-wrj8hrxp9Q" colab={"base_uri": "https://localhost:8080/"} outputId="423bda15-8805-4037-e7b6-65a7575a7933"
s = Sampler(small_ds,3,False)
[o for o in s]

# + id="cAyM-UXDxp9R" colab={"base_uri": "https://localhost:8080/"} outputId="8f96ed10-0c06-4f7a-ced1-08591b7ad9ff"
s = Sampler(small_ds,3,True)
[o for o in s]


# + id="Xf0r66xbxp9R"
def collate(b):
    xs,ys = zip(*b)
    return torch.stack(xs),torch.stack(ys)

class DataLoader():
    def __init__(self, ds, sampler, collate_fn=collate):
        self.ds,self.sampler,self.collate_fn = ds,sampler,collate_fn
        
    def __iter__(self):
        for s in self.sampler: yield self.collate_fn([self.ds[i] for i in s])


# + id="JfhQOgsVxp9S"
train_samp = Sampler(train_ds, bs, shuffle=True)
valid_samp = Sampler(valid_ds, bs, shuffle=False)

# + id="GZpHj2oqxp9S"
train_dl = DataLoader(train_ds, sampler=train_samp, collate_fn=collate)
valid_dl = DataLoader(valid_ds, sampler=valid_samp, collate_fn=collate)

# + id="qG_iithdxp9S" colab={"base_uri": "https://localhost:8080/", "height": 283} outputId="6367fbae-db63-4292-8db2-21ecef0bbbe4"
xb,yb = next(iter(valid_dl))
plt.imshow(xb[0].view(28,28))
yb[0]

# + id="wDtE1Hlgxp9S" colab={"base_uri": "https://localhost:8080/", "height": 283} outputId="f64d8585-9b27-4833-915c-3569b6e0bc0c"
xb,yb = next(iter(train_dl))
plt.imshow(xb[0].view(28,28))
yb[0]

# + id="fkS0eizexp9T" colab={"base_uri": "https://localhost:8080/", "height": 283} outputId="51ad1749-fbb5-42da-adb4-2f649e7adecc"
xb,yb = next(iter(train_dl))
plt.imshow(xb[0].view(28,28))
yb[0]

# + id="pwxmZdIBxp9U" colab={"base_uri": "https://localhost:8080/", "height": 449} outputId="975b6c12-4885-496a-9318-90d8529be2e9"
model,opt = get_model()
fit()

loss,acc = loss_func(model(xb), yb), accuracy(model(xb), yb)
assert acc>0.7
loss,acc

# + [markdown] id="BkOcYVz5xp9U"
# ### PyTorch DataLoader

# + [markdown] id="Ek63kxvpxp9U"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=4171)

# + id="8heUT6Qaxp9V"
#export
from torch.utils.data import DataLoader, SequentialSampler, RandomSampler

# + id="Mc6PTXcIxp9V"
train_dl = DataLoader(train_ds, bs, sampler=RandomSampler(train_ds), collate_fn=collate)
valid_dl = DataLoader(valid_ds, bs, sampler=SequentialSampler(valid_ds), collate_fn=collate)

# + id="URCJObCIxp9V" colab={"base_uri": "https://localhost:8080/", "height": 413} outputId="b316b685-e677-440a-a87a-179efba49375"
model,opt = get_model()
fit()
loss_func(model(xb), yb), accuracy(model(xb), yb)

# + [markdown] id="RV_tHff0xp9W"
# PyTorch's defaults work fine for most things however:

# + id="m1-XnjlYxp9W"
train_dl = DataLoader(train_ds, bs, shuffle=True, drop_last=True)
valid_dl = DataLoader(valid_ds, bs, shuffle=False)

# + id="hldX-jfTxp9W" colab={"base_uri": "https://localhost:8080/", "height": 449} outputId="1e61d521-5d11-4296-f054-d4b6da994933"
model,opt = get_model()
fit()

loss,acc = loss_func(model(xb), yb), accuracy(model(xb), yb)
assert acc>0.7
loss,acc


# + [markdown] id="4wHMsZhExp9W"
# Note that PyTorch's `DataLoader`, if you pass `num_workers`, will use multiple threads to call your `Dataset`.

# + [markdown] id="IZdHSkBbxp9X"
# ## Validation

# + [markdown] id="KzRnuK0dxp9X"
# You **always** should also have a [validation set](http://www.fast.ai/2017/11/13/validation-sets/), in order to identify if you are overfitting.
#
# We will calculate and print the validation loss at the end of each epoch.
#
# (Note that we always call `model.train()` before training, and `model.eval()` before inference, because these are used by layers such as `nn.BatchNorm2d` and `nn.Dropout` to ensure appropriate behaviour for these different phases.)

# + [markdown] id="gG-PcbGPxp9X"
# [Jump_to lesson 9 video](https://course.fast.ai/videos/?lesson=9&t=4260)

# + id="gZpIe-Jqxp9X"
def fit(epochs, model, loss_func, opt, train_dl, valid_dl):
    for epoch in range(epochs):
        # Handle batchnorm / dropout
        model.train()
#         print(model.training)
        for xb,yb in train_dl:
            loss = loss_func(model(xb), yb.long())
            loss.backward()
            opt.step()
            opt.zero_grad()

        model.eval()
#         print(model.training)
        with torch.no_grad():
            tot_loss,tot_acc = 0.,0.
            for xb,yb in valid_dl:
                pred = model(xb)
                tot_loss += loss_func(pred, yb.long())
                tot_acc  += accuracy (pred,yb)
        nv = len(valid_dl)
        print(epoch, tot_loss/nv, tot_acc/nv)
    return tot_loss/nv, tot_acc/nv


# + [markdown] id="ZU-f2jwTxp9Y"
# *Question*: Are these validation results correct if batch size varies?

# + [markdown] id="kavB4SYyxp9Y"
# `get_dls` returns dataloaders for the training and validation sets:

# + id="Ze8BCcJrxp9Z"
#export
def get_dls(train_ds, valid_ds, bs, **kwargs):
    return (DataLoader(train_ds, batch_size=bs, shuffle=True, **kwargs),
            DataLoader(valid_ds, batch_size=bs*2, **kwargs))


# + [markdown] id="gLfgd85jxp9Z"
# Now, our whole process of obtaining the data loaders and fitting the model can be run in 3 lines of code:

# + id="1xdz5p95xp9a" colab={"base_uri": "https://localhost:8080/"} outputId="0b4b7c1d-dcf1-4070-f5a1-8c9325b72364"
train_dl,valid_dl = get_dls(train_ds, valid_ds, bs)
model,opt = get_model()
loss,acc = fit(5, model, loss_func, opt, train_dl, valid_dl)

# + id="TM7lifJHxp9a"
assert acc>0.9

# + [markdown] id="I-fmPWSBxp9b"
# ## Export

# + id="LWl4Vw_Txp9b" outputId="21bd0810-c37a-4d12-9157-d733a96c0e26"
# !python notebook2script.py 03_minibatch_training.ipynb

# + id="OxGTEBHtxp9c"


StarcoderdataJupyterScriptsDedupFiltered