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

# + [markdown] id="HuHBa1Shto5L" colab_type="text"
# # Working with WandB and PyTorch

# + [markdown] id="1K3HBWwnvFy7" colab_type="text"
# ## Initialize it

# + id="cBxFPmNitopd" colab_type="code" colab={}
# !pip install wandb -q

# + id="No47NlZ4tlj0" colab_type="code" colab={}
from __future__ import print_function
import argparse
import random # to set the python random seed
import numpy # to set the numpy random seed
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
# Ignore excessive warnings
import logging
logging.propagate = False 
logging.getLogger().setLevel(logging.ERROR)

# WandB – Import the wandb library
import wandb


# + id="c5HyxwXmtwNb" colab_type="code" colab={}
# WandB – Login to your wandb account so you can log all your metrics
# !wandb login

# + [markdown] id="5sf3H6zavDVF" colab_type="text"
# ## Create a simple NN

# + id="BZ5VTzcCuzih" colab_type="code" colab={}
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        
        # In our constructor, we define our neural network architecture that we'll use in the forward pass.
        # Conv2d() adds a convolution layer that generates 2 dimensional feature maps to learn different aspects of our image
        self.conv1 = nn.Conv2d(3, 6, kernel_size=5)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        
        # Linear(x,y) creates dense, fully connected layers with x inputs and y outputs
        # Linear layers simply output the dot product of our inputs and weights.
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # Here we feed the feature maps from the convolutional layers into a max_pool2d layer.
        # The max_pool2d layer reduces the size of the image representation our convolutional layers learnt,
        # and in doing so it reduces the number of parameters and computations the network needs to perform.
        # Finally we apply the relu activation function which gives us max(0, max_pool2d_output)
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        
        # Reshapes x into size (-1, 16 * 5 * 5) so we can feed the convolution layer outputs into our fully connected layer
        x = x.view(-1, 16 * 5 * 5)
        
        # We apply the relu activation function and dropout to the output of our fully connected layers
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        
        # Finally we apply the softmax function to squash the probabilities of each class (0-9) and ensure they add to 1.
        return F.log_softmax(x, dim=1)


# + id="40QT1YcJvOW7" colab_type="code" colab={}
def train(args, model, device, train_loader, optimizer, epoch):
    # Switch model to training mode. This is necessary for layers like dropout, batchnorm etc which behave differently in training and evaluation mode
    model.train()
    
    # We loop over the data iterator, and feed the inputs to the network and adjust the weights.
    for batch_idx, (data, target) in enumerate(train_loader):
        if batch_idx > 20:
          break
        # Load the input features and labels from the training dataset
        data, target = data.to(device), target.to(device)
        
        # Reset the gradients to 0 for all learnable weight parameters
        optimizer.zero_grad()
        
        # Forward pass: Pass image data from training dataset, make predictions about class image belongs to (0-9 in this case)
        output = model(data)
        
        # Define our loss function, and compute the loss
        loss = F.nll_loss(output, target)
        
        # Backward pass: compute the gradients of the loss w.r.t. the model's parameters
        loss.backward()
        
        # Update the neural network weights
        optimizer.step()


# + [markdown] id="JZgbzo6IgAZK" colab_type="text"
# ## Define the Evaluation Step
#
# Here we add a line of code to:
#
# *   **wandb.log()** – Log your metrics (accuracy, loss and epoch) and examples of images along with the predicted and true labels. This allows you to visualize your neural network's performance over time.

# + id="MUI5GslnvWKW" colab_type="code" colab={}
def test(args, model, device, test_loader, classes):
    # Switch model to evaluation mode. This is necessary for layers like dropout, batchnorm etc which behave differently in training and evaluation mode
    model.eval()
    test_loss = 0
    correct = 0

    example_images = []
    with torch.no_grad():
        for data, target in test_loader:
            # Load the input features and labels from the test dataset
            data, target = data.to(device), target.to(device)
            
            # Make predictions: Pass image data from test dataset, make predictions about class image belongs to (0-9 in this case)
            output = model(data)
            
            # Compute the loss sum up batch loss
            test_loss += F.nll_loss(output, target, reduction='sum').item()
            
            # Get the index of the max log-probability
            pred = output.max(1, keepdim=True)[1]
            correct += pred.eq(target.view_as(pred)).sum().item()

            example_images.append(wandb.Image(data[0], caption = "Pred: {} Truth: {}".format(classes[pred[0].item()], classes[target[0]])))
    
    wandb.log({
        "Examples": example_images,
        "Test Accuracy": 100. * correct / len(test_loader.dataset),
        "Test Loss": test_loss})



# + [markdown] id="jH8yXrpKhDip" colab_type="text"
# # Train, Edit, and Retrain
# Run wandb.init() each time you start a new run. For this tutorial we're logging results to an open, shared project called "[pytorch-intro](https://app.wandb.ai/wandb/pytorch-intro/)" so everyone's results go into the same project. (The entity refers to the team name this project is under. In this case I've set it up as my own team, "wandb").
#
# You can see other runs on the [project page](https://app.wandb.ai/wandb/pytorch-intro/).
#
# ## Initialize Hyperparameters
#
# Here we add a few lines of code to:
# *   **wandb.init()** – Initialize a new W&B run. Each run is single execution of the training script.
# *   **wandb.config** – Save all your hyperparameters in a config object. This lets you use our app to sort and compare your runs by hyperparameter values.
#
# We encourage you to tweak these and run this cell again to see if you can achieve improved model performance!
#
# ## Track Results
# *   **wandb.watch()** – Fetch all layer dimensions, gradients, model parameters and log them automatically to your dashboard.
# *   **wandb.save()** – Save the model checkpoint.

# + id="rsV8d2RbvteT" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 69} outputId="8c37dca9-56de-4a51-8957-28af3699c11a"
wandb.init(entity="wandb", project="pytorch-intro")
wandb.watch_called = True

# Set up the configs
config = wandb.config
config.batch_size = 4
config.test_batch_size = 10
config.epochs = 50
config.lr = 0.1
config.momentum = 0.9
config.no_cuda = True  # Disable CUDA Training
config.seed = 42
config.log_inteval = 10 # How many batches to wait before logging


# + id="4C-7TXhQ1lW1" colab_type="code" colab={}
def main():
    use_cuda = not config.no_cuda and torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
    # Set random seeds and deterministic pytorch for reproducibility
    # random.seed(config.seed)       # python random seed
    torch.manual_seed(config.seed) # pytorch random seed
    # numpy.random.seed(config.seed) # numpy random seed
    torch.backends.cudnn.deterministic = True

    transform = transforms.Compose([
                                     transforms.ToTensor(),
                                     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    train_loader = torch.utils.data.DataLoader(datasets.CIFAR10(root='./data', train=True, download=True, transform=transform, 
                                                                ), shuffle=True, batch_size=config.batch_size, **kwargs)
    
    test_loader = torch.utils.data.DataLoader(datasets.CIFAR10(root='./data', train=False,
                                             download=True, transform=transform), batch_size=config.test_batch_size,
                                             shuffle=False, **kwargs)

    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    
    # Initialize our model, recursively go over all modules and convert their parameters and buffers to CUDA tensors (if device is set to cuda)
    model = Net().to(device)
    optimizer = optim.SGD(model.parameters(), lr=config.lr, momentum=config.momentum)

    # WandB – wandb.watch() automatically fetches all layer dimensions, gradients, model parameters and logs them automatically to your dashboard.
    # Using log="all" log histograms of parameter values in addition to gradients
    wandb.watch(model, log="all")

    for epoch in range(1, config.epochs + 1):
        train(config, model, device, train_loader, optimizer, epoch)
        test(config, model, device, test_loader, classes)
        
    # WandB – Save the model checkpoint. This automatically saves a file to the cloud and associates it with the current run.
    torch.save(model.state_dict(), "model.pt")
    wandb.save('model.pt')



# + id="PRIpPKx-4oWO" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 126, "referenced_widgets": ["a6809f5dfd2a4d59b42c508cab9590d4", "250290bee6a544ebbe555280bb56a0e1", "7270672815a6415bbb58fac758792ff8", "e8f720376e8b4812af9f8099d772e548", "4ba46186945c45cb83071e468648073a", "f0d0740ed5e1416f9949815c7590be33", "827727374e654121b0aa7403f97b369e", "d0b48b247a884fea9b9da1297a90293a"]} outputId="be9c1c0e-81f5-4fbe-a3d6-5c0b1637ec5d"
main()

# + id="MlVItjlJ4pyb" colab_type="code" colab={}


StarcoderdataJupyterScriptsDedupFiltered