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

# + colab={} colab_type="code" id="hqnl0AKVXIA4"
import matplotlib.pyplot
import numpy.random
import torch.utils.data
import torchvision
from torch import Tensor
from torch.nn import Module
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor

# + [markdown] colab_type="text" id="mCt-frpFH8oN"
# # Tutorial 1a. Logistic Regression
#
# In the first tutorial, we are going to train a logistic regressor on the MNIST dataset of handwritten digits. Next, we will turn this logistic regressor into a non-linear convolutional network.

# + colab={} colab_type="code" id="dzfEE578uSNp"
if torch.cuda.is_available():
    DEVICE = "cuda:0"
else:
    DEVICE = "cpu"
# -

# ## Loading Datasets
#
# The following code will load the MNIST dataset. Run it and inspect some of the images and their labels to confirm they are correct.

# + colab={"base_uri": "https://localhost:8080/", "height": 349, "referenced_widgets": ["59a8e0ee48e14c9691224d013d6c46fe", "ee8b4317d860412fa60bca6d1d0c935a", "bcc108058dc4490986ec987c568b2034", "ab7e6d830cc74c14876a052480eb185a", "3a488dd9262d42ac9b119cd37e2e71c9", "229d85ab3245444c9ad5b57c6ec58688", "a9ccabebd2c6428594120631518c505c", "<KEY>", "41f12163cbd84febb7fade8126ed7d5e", "<KEY>", "8f7373518a7840b38d00de7608a9e218", "e92daa0c7043415b8463dffd654e93e5", "<KEY>", "<KEY>", "1a38c10f1d0e4f20a4958d453757f9e7", "70658583c62a40c28345d43ceb70810b", "<KEY>", "d677eff7d5a4494087a60ff1db4dab92", "200b60fab8514008a571603c6a2519f8", "34a38f1355e54a24ad0757e91e016ff4", "8f98aca9cfae490dab5dccf54b939ce0", "04f7f68d4f784cf19f942ee29c24ecc7", "<KEY>", "<KEY>", "e11d5108f5a74dd6b7143172c693dc6a", "590b36901152442088a4c5e05e781b3d", "2de5166b2a40479fb69c08e7fda49594", "655321e930014d03a986ab862d9c01a1", "<KEY>", "<KEY>", "<KEY>", "<KEY>"]} colab_type="code" executionInfo={"elapsed": 7621, "status": "ok", "timestamp": 1586535092902, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gjq-QU8evTujx-PnMoIketNeW4yB_kgdtnRJ03IxnY=s64", "userId": "12331399045734156147"}, "user_tz": 240} id="GX_ky5qC--uv" outputId="97a811cf-59f5-4bb5-be96-03e19190cd4b"
train_dataloader = DataLoader(
    dataset=MNIST(
        root="/tmp/mnist",
        train=True,
        transform=ToTensor(),
        download=True,
    ),
    batch_size=64,
    shuffle=True,
)

test_dataloader = DataLoader(
    dataset=MNIST(
        root="/tmp/mnist",
        train=False,
        transform=ToTensor(),
        download=True,
    ),
    batch_size=64,
    shuffle=False,
)

# + colab={"base_uri": "https://localhost:8080/", "height": 62} colab_type="code" executionInfo={"elapsed": 7559, "status": "ok", "timestamp": 1586535092902, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gjq-QU8evTujx-PnMoIketNeW4yB_kgdtnRJ03IxnY=s64", "userId": "12331399045734156147"}, "user_tz": 240} id="E8dy7-kBHwco" outputId="66a535b5-83bc-48e0-f94a-1a79de8c3e66"
image, target = [*test_dataloader][0]

random_index = numpy.random.randint(0, 64)

image, target = image[random_index, 0], target[random_index]

matplotlib.pyplot.imshow(
    image, 
    cmap="gray", 
    interpolation="nearest",
)

matplotlib.pyplot.title(f"target = {target}")
matplotlib.pyplot.axis("off")

# + [markdown] colab_type="text" id="aKbXgidhWr6L"
# Next, we create a PyTorch dataloader for the MNIST dataset.

# + [markdown] colab_type="text" id="liekFZzvYX9E"
# Next, implement a logistic regression model in PyTorch. Note that a logistic regressor uses a linear transformation of the input.
#
# -

# ## Exercise 1. Logistic Regression Module

# + colab={} colab_type="code" id="zEw5YorSYkWF"
# class LogisticRegression(Module):
#     def __init__(self, input_size: int, num_classes: int):
#         super().__init__()
# 
#         ###########################################################################
#         # TODO: Instantiate the layer here.                                       #
#         ###########################################################################
# 
#     def forward(self, x: Tensor) -> Tensor:
#         ###########################################################################
#         # TODO: Apply the layer to the input.                                     #
#         ###########################################################################
#         pass

# +
module = LogisticRegression(28 * 28, 10)

module = module.to(device=DEVICE)

# +
###########################################################################
# TODO: Create criterion and optimizer here.                              #
###########################################################################

criterion = 

optimizer = 
# -

# ## Training

# We will use the following generic training loop for a PyTorch model.

# + colab={} colab_type="code" id="qVyEKl3OZLJw"
# Train the model. If everything is correct, the loss should go below 0.45.

EPOCHS = 5

# Exponential moving average of the loss:
ema = None

for epoch in range(EPOCHS):
    for batch_index, (train_images, train_targets) in enumerate(train_dataloader):
        train_images = train_images.view(-1, 28 * 28).requires_grad_().to(device=DEVICE)

        test_targets = train_targets.to(device=DEVICE)

        # Clear gradients w.r.t. parameters
        optimizer.zero_grad()

        # Forward pass to get output/logits
        outputs = module(train_images)

        # Calculate Loss: softmax --> cross entropy loss
        loss = criterion(outputs, train_targets)

        # Getting gradients w.r.t. parameters
        loss.backward()

        # Updates parameters:
        optimizer.step()

        # NOTE: It is important to call .item() on the loss before summing.
        if ema is None:
            ema = loss.item()
        else:
            ema += (loss.item() - ema) * 0.01

        if batch_index % 500 == 0:
            print(
                "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
                    epoch,
                    batch_index * len(train_images),
                    len(train_dataloader.dataset),
                    100.0 * batch_index / len(train_dataloader),
                    ema,
                ),
            )

# + [markdown] colab_type="text" id="oN8qTuGkZ-3X"
# **Question:** For the model you are currently using, is there any difference between using the model in `train` mode or using it in `eval` mode? 
#
# Create an SGD optimizer and us it to train the logistic regressor on the MNIST training data for a few epochs. What loss function do you need to use?

# + [markdown] colab_type="text" id="x0QN7Mhiar7d"
# ### Embeddings 
#
# Visualize the weights of the trained model. What do you see? Why?
# -

assert module.y.weight.shape == (10, 28 * 28)

# + colab={"base_uri": "https://localhost:8080/", "height": 167} colab_type="code" executionInfo={"elapsed": 37796, "status": "ok", "timestamp": 1586535123165, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gjq-QU8evTujx-PnMoIketNeW4yB_kgdtnRJ03IxnY=s64", "userId": "12331399045734156147"}, "user_tz": 240} id="gxletH44a4MX" outputId="1ba6c589-bcbb-4b86-a9ab-7fd4bc5bdfa1"
matplotlib.pyplot.imshow(
    numpy.transpose(
        torchvision.utils.make_grid(
            module.y.weight.view(10, 1, 28, 28),
            normalize=True,
            nrow=5,
        ), 
        (1, 2, 0),
    ), 
    interpolation="nearest",
)

matplotlib.pyplot.grid(False)
matplotlib.pyplot.gca().axis("off")
# -

# ## Evaluation

# + [markdown] colab_type="text" id="Lxf5NxQ6a5cT"
# Use the following function to measure the test accuracy of your trained model.

# + colab={} colab_type="code" id="8UYWbqZYa9Qr"
correct_predictions = 0
predictions = 0

# Iterate through test dataset
for test_images, test_targets in test_dataloader:
    test_images = test_images.view(-1, 28 * 28).to(DEVICE)

    # Forward pass only to get logits/output
    outputs = module(test_images)

    # Get predictions from the maximum value
    _, predicted = torch.max(outputs.data, 1)

    predictions += test_targets.size(0)

    if torch.cuda.is_available():
        correct_predictions += (predicted.cpu() == test_targets.cpu()).sum()
    else:
        correct_predictions += (predicted == test_targets).sum()

correct_predictions.item() / predictions

StarcoderdataJupyterScriptsDedupFiltered