# ---
# 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" colab={"base_uri": "https://localhost:8080/"} outputId="487bc35a-990c-4ca9-8de8-9f72d8d0370c"
# %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="db2c430b-a12a-4388-fd50-3c1f4621647f"
#Q1. Execute the following statement. What is displayed? What does it mean?
# !pwd

# + id="jhAroxykqtRF" colab={"base_uri": "https://localhost:8080/"} outputId="8d33081e-b595-456b-ea92-a8162ba21588"
#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="I86djNvcxp70"
#export
#Q3a: Execute the following statement. Read the error message. Explain what it means. 

from exp.nb_02 import *


# + id="KsprL6o5IRIe"
#Q3b. You can solve the problem by executing the following statement before the above statement. 
# Explain why the following statements can solve this "module not found" problem. 

# + id="0viyrQOz8TVm"
import sys
sys.path.append('/content/course-v3/nbs/dl2')
sys.path

# + id="jDlj09yaftUr"
import torch.nn.functional as F

# + id="FOZb4AeKq6vQ"


# + [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"
#Q4: 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()

# + id="oHcenWs_rZ_I" colab={"base_uri": "https://localhost:8080/"} outputId="cbed36bb-279e-44f1-8c91-f887f63918ff"
#Q5a: Execute the following statement. What does the number displayed mean?
len(x_train)

# + id="mxSwa7g5rcer"
#Q5b: Execute the following statement. What do the numbers displayed refer to?
x_train[0]

# + colab={"base_uri": "https://localhost:8080/"} id="PTrGABBgiPvi" outputId="3b39d99a-6a71-4059-f0ea-039d08181b11"
#Q5c: Execute the following statement. What do the numbers displayed refer to?
y_train[0]

# + colab={"base_uri": "https://localhost:8080/"} id="pzZC2MDkhtdJ" outputId="28d35fd5-a0b4-442f-a8d4-9c4ffc6c84c0"
#Q5d. Execute the following statement. What does the number displayed refer to? 
x_train[0].shape

# + id="pA-uJSgpD1Z0" colab={"base_uri": "https://localhost:8080/"} outputId="f5e1dda6-060e-46fe-dcda-1f6b2612339a"
#Q5e. Execute the following statement. What do the numbers displayed refer to?
x_train.shape

# + colab={"base_uri": "https://localhost:8080/"} id="RkNGxyrOibgh" outputId="6a14de35-419b-4027-d0ee-29ca69dd1cc5"
#Q5f. Execute the following statement. What do the numbers displayed refer to?
y_train.shape

# + id="DESekbEYxp77"
#Q6: 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"
#The following defines Model class, which will be used to create a neural net.
#Q7a: nn.Module is the parent class of class Model. Why do you want to make Model a child class to its parent nn.Module, rather than making Model stand on its own?
#Q7b. See the definition of self.layers field. It contains nn.Linear(n_in,nh). What is the difference between nn.Linear class nn.Linear(n_in,nh)? 
#Q7c. Would you think that the weight and bias parameters of the two linear layers are initialized when object nn.Linear(n_in,nh) is constructed? If so, guess why.

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"
#Q8. When the following statement is executed, what happens? Explain by referring to a method in class Model.
model = Model(m, nh, 10)

# + id="gPFt346vxp78"
#Q9a. When the following statement is executed, what happens? Explain by referring to a method in class Model. 
#Q9b. Draw a diagram or a graph by using arrows and boxes, which shows what computation is performed when model(x_train) is executed.
pred = model(x_train)

# + id="gXbLts9ZRMgd" colab={"base_uri": "https://localhost:8080/"} outputId="431ce4ba-3128-4584-d0d1-ee292ff6c87c"
#Q10: Execute the following statement. What do the 10 displayed numbers refer to?
pred[0]

# + id="-Ejrm01LRSxj" colab={"base_uri": "https://localhost:8080/"} outputId="532212cf-bef6-4e34-80a4-4f43e2b6acde"
#Q11. Execute the following statement. Is the resulting value near to 1? If not, what does it imply?
pred[0].sum()


# + [markdown] id="qhbGAsGWxp79"
# ### Softmax function

# + [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 softmax(x): return (x.exp()/(x.exp().sum(-1,keepdim=True)))


# + id="Wo_n_z36r77v"
sm_pred = softmax(pred)


# + id="cnURmMOHsA2T" colab={"base_uri": "https://localhost:8080/"} outputId="3f1f3bf2-7a74-418f-f4e6-6f245650a913"
#Q12. Execute the following statement. What do the 10 displayed numbers refer to?
sm_pred[0]

# + id="MvUVm4qpsQ80" colab={"base_uri": "https://localhost:8080/"} outputId="a6894356-2a01-4248-c2e5-45e4939cb8a3"
#Q13. Execute the following statement. What does the displayed numbers mean?
sm_pred[0].sum()


# + [markdown] id="sqq6KTqimJ-T"
#

# + [markdown] id="fWPxYHo1mNjm"
# ## Log softmax

# + id="DNuWma6zr00h"
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="7abf227a-2167-410e-f176-46381e1f7339"
log_sm_pred[0]

# + [markdown] id="DfNuX2v1xp8B"
#
#
# ## Cross Entropy Loss Function.
#  Read the following paragraph to understand what the cross entroy loss functioni is. 
#
# The difference between  two probabilities: 
#  https://datascience.stackexchange.com/questions/20296/cross-entropy-loss-explanation
#
# The cross entropy formula takes in two distributions,  $p(x^{(s)})$, the true distribution (defined by the label data), and $\hat{p}(x^{(s)})$, the estimated distribution (predicted by the neural net), defined over the discrete variable $x^{(s)}$  and is given by
#
# $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
#
#  \begin{equation}
#  H(p,\hat{p})= -\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)}) ) 
#   \tag{crossEntroyEq}
#  \end{equation}
#  
#  Te softmax function plays the role of the probability distribution $\hat{p}_{1} (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"
# # integer array indexing 

# + colab={"base_uri": "https://localhost:8080/"} id="VehY657jB7IB" outputId="756e9a60-b726-4cee-f668-1d9629ab1d2a"
#Q14a.  Execute the following statement. What is the role of list [0,1,2] in the statement?  What do the displayed numbers refer to? 
log_sm_pred[ [0,1,2]]

# + id="uJM_9p0rxp8H" colab={"base_uri": "https://localhost:8080/"} outputId="2a209d3b-eab0-485b-e242-ea4beeb58c28"
#Q14b.  Execute the following statement. What is the role of list [2,4,6] in the statement?  What do the displayed numbers mean? Explain compare the result of Q15b and the result of this statement.
log_sm_pred[[0,1,2], [2,4,6]]

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

# + id="xoc2YwU5wATE" colab={"base_uri": "https://localhost:8080/"} outputId="3394e65f-8e87-47fc-9966-d0c0ddf90987"
range(y_train.long().shape[0])

# + id="va2lulaep1_J"


# + id="8GMBhmq9w4QG"
selectOneHot = log_sm_pred[ range(y_train.long().shape[0]), y_train.long() ]

# + colab={"base_uri": "https://localhost:8080/"} id="v1c7b4HtqWLi" outputId="7ce32728-0fc7-4f5c-c14b-990846c8fd55"
selectOneHot.shape


# + id="1ZDSRenHqdcP"
#Q15.  What does selectOneHot refer to?

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


# + id="vzmLXXWiPTSF" colab={"base_uri": "https://localhost:8080/"} outputId="cc1d6106-f6c5-4bbe-a838-58c6eb6e0653"
y_train.long().shape[0]

# + id="Bs5UU4fyxp8K"
#Q16. Read nill function and explain how this function computes the result. You need to refer to {CrossEntropyLoss}.
loss = nll(log_sm_pred, y_train.long())

# + id="zkSkgQYQxp8L" colab={"base_uri": "https://localhost:8080/"} outputId="6c5cbca1-0047-4d62-b275-7445ec78e38e"
loss

# + id="-Xb-DbeRq2aa"
#Q16. Compare function nll() and the formula CrossEntropyEq. Explain that function nll() computes the cross entropy function between
# predicted probability distrubution of the input images and the ground truth (labeled) probability distribution of the input images. 

# + [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 built-in Pytorch function  combines log_softmax and nll in a single function.
# In the following code, we will use this function, not the function nll defined above.

# + 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="ad677304-2fd7-412f-8184-a78036b63f0d"
bs=1000                  # batch size

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

#Q17. What does preds[0] refer to?

# + id="Ue2zaDOLxp8a" colab={"base_uri": "https://localhost:8080/"} outputId="d5a5fd32-2d14-4f32-f608-76495eb48a88"
yb = y_train[0:bs]
loss = loss_func(preds, yb.long())
loss
#Q18. What does the value of loss refer to?

# + id="vYSsK_Rlxp8b" colab={"base_uri": "https://localhost:8080/"} outputId="dd10d65f-cd79-4188-c309-151f3085a488"
accuracy(preds, yb)
#Q19. What does  the value of the above statement refer to? Examine the value. Is it supposed to be equal to one approximately?
# If it is not at the current moment, what is the reason for it?

# + [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
#

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

# + id="Q0KH3m-Uxp8c"
#Q20. What does the number  (n-1)//bs + 1 refer to? 
import pdb
for epoch in range(epochs):
    for i in range((n-1)//bs + 1): # for each batch in the current epoch
#       pdb.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())        

        #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                
        #The following performs what Pytorch function optimizer.step() does:
        # It updates all the weight and bias parameters based on the gradients of loss with respect to the weight and bias parameter 
        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

                    #
                    # loss.backward() computes the gradient of loss w.r.t. graph LEAVES.
                    #  This function accumulates gradients in the leaves - you might need to zero
                    # 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_()

       #Q21:  print the loss and the accuracy after training the net using the validation dataset.
       #      Explain how the loss and accuracy change as each batch is used  for training the network model.
       #      Try to convey detailed and specific information about the progress of training the neural net.
       #      Observe the printed data carefully
    
       #Q22. Explain why you would use validation dataset to check the progress of learning of your network.
                   
        
        yHatValid = model(x_valid)
        lossValid = loss_func( yHatValid, y_valid.long() )        

        
        print('epoch={0}, batch ={1}:'.format(epoch, i) )  
        print('  lossValid=', lossValid )
        print('  accuracyValid = ', accuracy( yHatValid, y_valid) )

    #Q23. Afer each epoch, print the loss and the accuracy of the network by using the training dataset. 
    #     Explain the result. Be attentive to the result.   

    yHatTrain = model(x_train)
    lossTrain = loss_func( yHatTrain, y_train.long() )        

    print('*******************************************************')    
    print('epoch={0}:', epoch )  
    print('  lossTrain=', lossTrain)
    print('  accuracyTrain = ', accuracy( yHatTrain, y_train) )

    #Q24. Afer each epoch, print the loss and the accuracy of the network by using the validatio  dataset. 
    #     Explain the result. Be attentive to the result.   

    yHatValid = model(x_valid)
    lossValid = loss_func( yHatValid, y_valid.long() )        

    print('*******************************************************')    
    print('epoch={0}:', epoch )  
    print('  lossValid ', lossValid)
    print('  accuracyValid = ', accuracy( yHatValid, y_valid) )
        

# + [markdown] id="98nX1aGLyXJP"
#

# + id="I0fC9kbvxp8g" colab={"base_uri": "https://localhost:8080/"} outputId="8cd813d0-b5cb-4df3-97c0-bf7ae1529cfb"
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="56949b22-3eae-4ec0-8a23-e2f754e7fa66"
for name,l in model.named_children(): print(f"{name}: {l}")

# + id="QQAcSLAVxp8l" colab={"base_uri": "https://localhost:8080/"} outputId="3607efe6-16e9-4fb8-e118-a6d7db48cb32"
model

# + id="O77Ds-LRxp8m" colab={"base_uri": "https://localhost:8080/"} outputId="8dd994eb-d95f-44f3-9adb-24ec6f4bb3d1"
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="e98cf04b-c832-4d83-b86f-630125c2084a"
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="40d56eb6-6822-436c-c011-bd7ae14336d3"
mdl = DummyModule(m,nh,10)
mdl

# + id="asWZtQT7xp8r" colab={"base_uri": "https://localhost:8080/"} outputId="ba4de7ed-d661-41a1-c769-f9dbade23efa"
[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="8148eec4-6e20-4ff3-8949-04c92b463be5"
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="fd6d8b69-bbf4-40a8-c02e-54b28a92cd6e"
model

# + id="QheaMa-6xp8y" colab={"base_uri": "https://localhost:8080/"} outputId="394de92e-5802-4616-9845-97c940736af1"
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="3afc0d73-3d2e-4fb1-e601-c8dcbdff0bfb"
fit()
loss_func(model(xb), yb.long()), accuracy(model(xb), yb)

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

# + id="s1EvUMznxp81" colab={"base_uri": "https://localhost:8080/"} outputId="bc2353bd-57a8-4ba3-84fb-95360c0575ce"
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="5486e4ff-e6d2-48d8-fcb7-0221c66b98fa"
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="647cfcd2-e327-412e-f4f6-21373479b771"
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="1b67c1be-aac4-44b2-fc06-c3ca97a47774"
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="c7e66198-b65c-4bf5-d0bc-186c336644a3"
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="c824b5eb-dfbd-443e-aa68-8ff518caf9aa"
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="36a6bc37-1a9d-45e5-84ea-91ea3ca2f72c"
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="be1bcbd7-fe53-4bff-f68f-484e17a56168"
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="16c68025-e171-44dd-e22f-c07442e89e3f"
s = Sampler(small_ds,3,False)
[o for o in s]

# + id="cAyM-UXDxp9R" colab={"base_uri": "https://localhost:8080/"} outputId="e16e8f20-47c9-4e77-877d-50a665aee1c5"
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="2ca0cbc3-b121-4a71-8ede-aa4b3a085f4c"
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="1afcd8fd-2ba0-4474-bafd-2b4f7400b85e"
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="779a6399-14ee-464f-e501-7b00dfcd6816"
xb,yb = next(iter(train_dl))
plt.imshow(xb[0].view(28,28))
yb[0]

# + id="pwxmZdIBxp9U" colab={"base_uri": "https://localhost:8080/", "height": 413} outputId="c186f148-4b56-4150-99a1-e2bc5b92efcf"
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"
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"
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"
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"
# !python notebook2script.py 03_minibatch_training.ipynb

# + id="OxGTEBHtxp9c"


StarcoderdataJupyterScriptsDedupFiltered