Introduction to PyTorch

import torch # import pytorch

z = torch.zeros(5,3)
print(z)
print(z.dtype)

i = torch.ones((5,3),dtype=torch.int16) # Overide default dtype
print(i.dtype)

torch.manual_seed(42) # Seed PRNG
rn = torch.rand(2,2)
print("A random tensor:",rn)

ones = torch.ones((2,3))
print(ones)
twos = torch.ones((2,3),dtype=torch.int16)*2
print(twos)
print(ones+twos)

r = (torch.rand((2,2)) - 0.5) * 2 # Values between -1 and 1
print("A random matrix:")
print(r)

# Common Mathematical operations are supported
print("\nAbosulte value of r:")
print(torch.abs(r))

# As are trigonometric functions:
print("\nInverse sine of r:")
print(torch.asin(r))

# and linear algrebra operations like determinant and singular value decomposition
print("\nDeterminant of r:")
print(torch.det(r))
print("\nSingular Value Decomposition of r:")
print(torch.svd(r))

# and statsistical and aggegrate operations
print("\nAverage and Standard Deviation of r:")
print(torch.std_mean(r))
print("\nMax value of r:")
print(torch.max(r))

import torch # for all things PyTorch
import torch.nn as nn # for torch.nn.Module, the parent object for PyTorch modules
import torch.nn.functional as F # for the activation fucntion

class LeNet(nn.Module):
  """
  Implementation of an abbridged version of LeNet-5
  Typical PyTprch Model:
  - Inherits from torch.nn.Module
  - A model will have an __init__() function, where it instantiates its layers, and loads any artifacts it might need (e.g., an NLP might need a vocab)
  - A model will have a forward() function. This is where the actual computation happens. An input is passed through the network layers and various functions to generate an output
  - Other than that, you can build a model however you want
  """
  def __init__(self):
    super(LeNet,self).__init__()
    # 1 input image channel (black and white), 6 output channels, 5x5 square convolution
    # kernel
    self.conv1 = nn.Conv2d(1,6,5)
    self.conv2 = nn.Conv2d(6,16,5)
    # an affine operation: y = Wx+b
    self.fc1 = nn.Linear(16*5*5,120) # 5*5 for image dimension
    self.fc2 = nn.Linear(120,84)
    self.fc3 = nn.Linear(84,10)
  def forward(self,x):
    # Max Pooling over a (2,2) window
    x = F.max_pool2d(F.relu(self.conv1(x)),(2,2))
    # If the size is a square you can only specify a single number
    x = F.max_pool2d(F.relu(self.conv2(x)),2)
    x = x.view(-1, self.num_flat_features(x))
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.fc3(x)
    return x
  def num_flat_features(selx,x):
    size = x.size()[1:] # all dimensions except the batch dimension
    num_features = 1
    for s in size:
      num_features *= s
    return num_features

net = LeNet()
print(net) # What does the object tell us about itself?
"""
A subclass of torch.nn.Module will report the layers it has created and their shapes and parameters. This can provide a handy overview of a model if you wnat to get the gist of its processing.
"""

input = torch.rand(1,1,32,32) # stand-in for a 32x32 black and white image

print('\nImage batch shape:')
print(input.shape)
"""
PyTorch models assume they are working with batches of data
"""

output = net(input) # We don;t call forward() directly
"""
Asking the model for an inference. The output of this call represents the model's confidence that the input represents a particualr digit.
"""
print('\nRaw output:')
print(output)
print(output.shape)

%matplotlib inline
import torch
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose([
    transforms.ToTensor(), # Converts images loaded by Pillow to PyTorch tensors
    # Adjusts the values of the tensor so that their avaerage is zero
    # and their standard deviation is 1.0. Most activation functions have their
    # strongest gradient around 0, so centering data their can speed learning
    # The values passed to the transform are the means (first tuple) and std
    # deviations of the rgb images in the dataset.
    transforms.Normalize(
        (0.4914, 0.4822, 0.4465),
         (0.2470, 0.2435, 0.2616)
    )
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)

"""
Give batches of 4 images from trainset, randomize order, and use two
workers to load data
"""
trainloader = torch.utils.data.DataLoader(
    trainset,
    batch_size=4, # 4 images from trainset
    shuffle=True, # randomize order
    num_workers=2 # use two workers to load data
)

import matplotlib.pyplot as plt
import numpy as np

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))


# get some random training images
dataiter = iter(trainloader)
images, labels = next(dataiter)

# show images
imshow(torchvision.utils.make_grid(images))
# print labels
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))

#%matplotlib inline

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import torchvision
import torchvision.transforms as transforms

import matplotlib
import matplotlib.pyplot as plt
import numpy as np

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

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

import matplotlib.pyplot as plt
import numpy as np

# functions to show an image


def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))


# get some random training images
dataiter = iter(trainloader)
images, labels = next(dataiter)

# show images
imshow(torchvision.utils.make_grid(images))
# print labels
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()

criterion = nn.CrossEntropyLoss() # Loss Function
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) # Optimizer

for epoch in range(2):  # loop over the dataset multiple times

    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # get the inputs
        inputs, labels = data

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs) # Ask the model for its predictions
        loss = criterion(outputs, labels) # Compute difference between correct and predicted outputs
        loss.backward() # Calculate the gradients that will direct the learning
        optimizer.step() # Perform one leanring step
        # uses the gradients from teh backward() call to nudge the
        # learning weights in the direction that it thinks will reduce the loss

        # print statistics
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

print('Finished Training')

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (
    100 * correct / total))

import torch
import math
"""
Creating Tensors
"""
# empty() call allocates memort for the tensor, but does not initialize it with
# any values
x = torch.empty((3,4)) # Simplest way to create a tensor
print(x)

zeros = torch.zeros((2,3))
print(zeros)

ones = torch.ones((2,3))
print(ones)

torch.manual_seed(1729)
random = torch.randn((2,3))
print(random)

"""
Tensor Shapes
"""
x = torch.empty((2,2,3))
print(x.shape)
empty_like_x = torch.empty_like(x)
print(empty_like_x.shape)
zeros_like_x = torch.zeros_like(x)
print(zeros_like_x.shape)
ones_like_x = torch.ones_like(x)
print(ones_like_x.shape)
random_like_x = torch.randn_like(x)
print(random_like_x.shape)

"""
Tensor Data Types
"""
a = torch.ones((2,3),dtype=torch.int16)
b = torch.randn((2,3),dtype=torch.float64) * 20
print(b)
c = b.to(torch.int32)
print(c)

"""
Math and Logic with PyTorch Tensors
"""
ones = torch.zeros((2,2)) + 1
twos = ones * 2
threes = (torch.ones((2,2)) * 7 - 1) / 2
fours = twos ** 2
sqrt2s = twos ** 0.5

print(ones)
print(twos)
print(threes)
print(fours )
print(sqrt2s)

powers2 = twos ** torch.tensor([[1,2],[3,4]])
print(powers2)
fives = ones + fours
print(fives)
dozens = threes * fours
print(dozens)

"""
More Math with Tensors
"""
a = torch.rand((2,4)) * 2 - 1
# Common Functions
print(torch.clamp(a,-0.5,0.5))
# Comparisons
d = torch.tensor([[1.,2.],[3.,4.]])
e = torch.ones(1,2) # Many comparison operations support broadcasting
print(torch.eq(d,e)) # Returns a Tensor of type bool
# Reductions
print(torch.max(d)) # Returns a single-element tensor
print(torch.max(d).item()) # Extracts the value from the returned tensor
print(torch.unique(torch.tensor([1,2,1,2,1,2,3]))) # Filter unique elements

# vector and linear algebra operations
v1 = torch.tensor([1., 0., 0.])         # x unit vector
v2 = torch.tensor([0., 1., 0.])         # y unit vector
m1 = torch.rand(2, 2)                   # random matrix
m2 = torch.tensor([[3., 0.], [0., 3.]]) # three times identity matrix
print(torch.linalg.cross(v2, v1)) # negative of z unit vector (v1 x v2 == -v2 x v1)
print(m1)
m3 = torch.linalg.matmul(m1, m2)
print(m3)                  # 3 times m1
print(torch.linalg.svd(m3))       # singular value decomposition


# Altering Tensor in Place
b = torch.tensor([0, math.pi / 4, math.pi / 2, 3 * math.pi / 4])
print('\nb:')
print(b)
print(torch.sin_(b))  # note the underscore
print(b)

a = torch.ones(2, 2)
b = torch.rand(2, 2)

print('Before:')
print(a)
print(b)
print('\nAfter adding:')
print(a.add_(b))
print(a)
print(b)

"""
Changing a (3,226,226) tensor to (1,2,226,226) tensor
"""
a = torch.rand(3, 226, 226)
"""
The unsqueeze() method adds a dimension of extent 1. unsqueeze(0) adds it
as a new zeroth dimension - so now you have a batch of one.
"""
b = a.unsqueeze(0)

print(a.shape)
print(b.shape)

a = torch.rand(1, 20)
print(a.shape)
print(a)

b = a.squeeze(0)
print(b.shape)
print(b)

c = torch.rand(2, 2)
print(c.shape)

d = c.squeeze(0)
print(d.shape)

# %matplotlib inline

import torch

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import math

# Setting requires_grad = True means that in every computation that
# follows, autograd will be accumulating the history of the computation in
# the output tensors of that computation.
a = torch.linspace(0., 2. * math.pi, steps=25, requires_grad=True)
print(a)
b = torch.sin(a)
plt.plot(a.detach(), b.detach(),label=r"2 $ \cdot $ sin(a)")


c = 2*b
d = c+1
out = d.sum()

"""
Each grad_fn stored with our tensors allows you to walk the computation
all the way back to its inputs with its next_function property. We can
see below that drilling down this property on d dhows us the
gradient functions for all prior tensors.
"""
print('d:')
print(d.grad_fn)
print(d.grad_fn.next_functions)
print(d.grad_fn.next_functions[0][0].next_functions)
print(d.grad_fn.next_functions[0][0].next_functions[0][0].next_functions)
print(d.grad_fn.next_functions[0][0].next_functions[0][0].next_functions[0][0].next_functions)
print('\nc:')
print(c.grad_fn)
print('\nb:')
print(b.grad_fn)
print('\na:')
print(a.grad_fn)

out.backward()
print(a.grad)

plt.gcf().gca().plot(a.detach(), a.grad.detach(),'-r',label=r"2 $ \cdot $ cos(a)")
plt.legend()
plt.show()

import torch

class TinyModel(torch.nn.Module):
  """
  Example of a very simple model with two linear layers and an activation
  function
  """
  def __init__(self):
    super(TinyModel,self).__init__()
    self.linear1 = torch.nn.Linear(100,200)
    self.activation = torch.nn.ReLU()
    self.linear2 = torch.nn.Linear(200,10)
    self.softmax = torch.nn.Softmax()
  def forward(self,x):
    x = self.linear1(x)
    x = self.activation(x)
    x = self.linear2(x)
    x = self.softmax(x)
    return x

tinymodel = TinyModel()
print("The model:")
print(tinymodel)

print("\n\nJust one Layer:")
print(tinymodel.linear2)

print('\n\nModel params:')
for param in tinymodel.parameters():
    print(param)

print('\n\nLayer params:')
for param in tinymodel.linear2.parameters():
    print(param)

"""
Tensor Initialization
"""
import torch
data = [[1,2], [3,4]]
x_data = torch.tensor(data)
import numpy as np
np_array = np.array(data)
x_np = torch.from_numpy(np_array)
x_ones = torch.ones_like(x_np)
x_rand = torch.rand_like(x_data,dtype=torch.float)
"""
With Random or Constant Values
"""
shape = (2,3,)
rand_tensor = torch.rand(shape)
ones_tensor = torch.ones(shape)
zeros_tensor = torch.zeros(shape)
"""
Tensor Attributes
"""
print(f"Shape of tensor: {rand_tensor.shape}")
print(f"Datatype of tensor: {rand_tensor.dtype}")
print(f"Device tensor is stored on: {rand_tensor.device}")
"""
Tensor Operations
"""
if torch.cuda.is_available():
  device = torch.device('cuda')
else:
  device = torch.device('cpu')

"""
Standard NumPy-like Indexing and Slicing
"""
tensor = torch.ones((4,4),device=device)
tensor[:,1] = 0
print(tensor)

"""
Joining Tnesors:

You can use torch.cat to concatenate a sequence of tensors along a given dimension
"""
t1 = torch.cat([tensor,tensor,tensor,tensor],dim=1)
print(t1)

"""
Multiplying Tensors
"""
# This computes the element-wise product
print(f"tensor.mul(tensor) \n {tensor.mul(tensor)} \n")
# Alternative syntax:
print(f"tensor * tensor \n {tensor * tensor}")

print(f"tensor.matmul(tensor.T) \n {tensor.matmul(tensor.T)} \n")
# Alternative syntax:
print(f"tensor @ tensor.T \n {tensor @ tensor.T}")

import torch
# Loas a pretrained resnet18 model from torchvision
from torchvision.models import resnet18, ResNet18_Weights
model = resnet18(weights=ResNet18_Weights.DEFAULT)
data - torch.rand((1,3,64,64)) # 3 - channel 64 x 64 image batch of 1
labels = torch.rand(1,1000) # Label initialized to some random values
prediction = model(data) # forward pass
loss = (prediction - labels).sum() # Calculate loss
# The backward() method kicks off backpropagation.
# Autograd calculates and stroes the gradientrs for each model parameter in the parameter's .grad attribute
loss.backward() # Backward pass

# Load an optimizer - in this case SGD with a learning rate of 0.01 and a momentum of 0.9
optim = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)
# Call step() to initiate gradient descent. The optimizer adjusts
# each parameter by its gradient stored in .grad
optim.step()

import math
import torch

class LegendrePolynomial3(torch.autograd.Function):
    """
    We can implement our own custom autograd Functions by subclassing
    torch.autograd.Function and implementing the forward and backward passes
    which operate on Tensors.
    """

    @staticmethod
    def forward(ctx, input):
        """
        In the forward pass we receive a Tensor containing the input and return
        a Tensor containing the output. ctx is a context object that can be used
        to stash information for backward computation. You can cache arbitrary
        objects for use in the backward pass using the ctx.save_for_backward method.
        """
        ctx.save_for_backward(input)
        return 0.5 * (5 * input ** 3 - 3 * input)

    @staticmethod
    def backward(ctx, grad_output):
        """
        In the backward pass we receive a Tensor containing the gradient of the loss with respect to the output, and we need to compute the gradient of
        the loss with respect to the input.
        """
        input, = ctx.saved_tensors
        return grad_output * 1.5 * (5 * input ** 2 - 1)

class Polynomial3(torch.nn.Module):
    def __init__(self):
        """
        In the constructor we instantiate four parameters and assign them as
        member parameters.
        """
        super().__init__()
        self.a = torch.nn.Parameter(torch.randn(()))
        self.b = torch.nn.Parameter(torch.randn(()))
        self.c = torch.nn.Parameter(torch.randn(()))
        self.d = torch.nn.Parameter(torch.randn(()))

    def forward(self, x):
        """
        In the forward function we accept a Tensor of input data and we must
        return a Tensor of output data. We can use Modules defined in the
        constructor as well as arbitrary operators on Tensors.
        """
        return self.a + self.b * x + self.c * x ** 2 + self.d * x ** 3

    def string(self):
        """
        Just like any class in Python, you can also define custom method on PyTorch modules
        """
        return f'y = {self.a.item()} + {self.b.item()} x + {self.c.item()} x^2 + {self.d.item()} x^3'

from pathlib import Path
import requests

DATA_PATH = Path("data")
PATH = DATA_PATH / "mnist"
PATH.mkdir(parents=True,exist_ok=True)

URL = "https://github.com/pytorch/tutorials/raw/main/_static/"
FILENAME = "mnist.pkl.gz"

if not (PATH / FILENAME).exists():
  content = requests.get(URL + FILENAME).content
  (PATH / FILENAME).open("wb").write(content)

"""
The dataset is in numpy array format, and has been stored using pickle
"""
import pickle
import gzip

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

"""
Each image is 28 x 28, and is being stored as flattened row of length 784.
"""

from matplotlib import pyplot
import numpy as np

pyplot.imshow(x_train[0].reshape((28, 28)), cmap="gray")
# ``pyplot.show()`` only if not on Colab
try:
    import google.colab
except ImportError:
    pyplot.show()
print("x_train shape:",x_train.shape)

"""
PyTorch uses torch.tensor, rather than numpy arrays, so we need to convert our
data
"""
import torch

x_train, y_train, x_valid, y_valid = map(
    torch.tensor, (x_train, y_train, x_valid, y_valid)
)
n, c = x_train.shape
print("x_train shape (after transformation):",x_train.shape)
print("y_train min, y_train max:",y_train.min(),",",y_train.max())

import math
# Xavier Initialization (http://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf)
weights = torch.randn(784,10) / math.sqrt(784)
# Setting requires_grad after initialization, since we don;t want that step
# included in the gradient
weights.requires_grad_()
bias = torch.zeros(10, requires_grad=True)

def log_softmax(x):
  """
  Activation function
  """
  return x - x.exp().sum(-1).log().unsqueeze(-1)

def model(xb):
  """
  Matrix multiplication and broadcasted addition
  to create a simple linear model
  """
  return log_softmax(xb @ weights + bias)

bs = 64  # batch size

xb = x_train[0:bs]  # a mini-batch from x
# Call the function on the data - forward pass
preds = model(xb)  # predictions
print(preds[0], preds.shape)

def nll(input, target):
  """
  Non negative log-likelihood to use as the loss function
  """
  return -input[range(target.shape[0]), target].mean()

loss_func = nll
yb = y_train[0:bs]
print(loss_func(preds, yb))

def accuracy(out, yb):
  """
  Function to calculate the accuracy of the model. For each
  prediction, if the index with the largest value matches the target
  value, then the prediction was correct
  """
  preds = torch.argmax(out, dim=1)
  return (preds == yb).float().mean()

print(accuracy(preds, yb)) # Accuracy of the model

# from IPython.core.debugger import set_trace
"""
Running a training loop.
"""

lr = 0.5  # learning rate
epochs = 2  # how many epochs to train for

for epoch in range(epochs):
  for i in range((n - 1) // bs + 1):
    #         set_trace()
    """
    Selecting a mini-batch of data
    """
    start_i = i * bs
    end_i = start_i + bs
    xb = x_train[start_i:end_i]
    yb = y_train[start_i:end_i]
    # Making predictions - the forward pass
    pred = model(xb)
    # Calculating the loss
    loss = loss_func(pred, yb)
    # Update the gradients of the model
    loss.backward()

    with torch.no_grad():
      weights -= weights.grad * lr # Updating weights
      bias -= bias.grad * lr # Updating biases
      """
      Setting weight gradient to zero so that we are ready for next training
      loop. Otherwise, the gradients would record a running tally of all the
      operations that has happened (loss.backward() adds the gradients to whatever is already stored, rather than replacing them).
      """
      weights.grad.zero_()
      bias.grad.zero_()

"""
Refactor code to take advantage of PyTorch;s nn classes to make the code more concise and flexible
"""
import torch.nn.functional as F
"""
Activation and loss functions can be found in the torch.nn.functional submodule (which is generally imported into namespace F by convention).

If you are using negative log likelihood and log softmax activation, then PyTorch provides a single function F.cross_entropy that combines the two.
"""
loss_func = F.cross_entropy
def model(xb):
  return xb @ weights + bias
print(loss_func(model(xb), yb), accuracy(model(xb), yb))

"""
Refactour using nn.Module and nn.Parameter, for a cleaner and more concise
training loop.
"""

from torch import nn

class Mnist_Logistic(nn.Module):
  def __init__(self):
    super().__init__()
    self.weights = nn.Parameter(torch.randn(784,10)  / math.sqrt(784) )
    self.bias = nn.Parameter(torch.zeros(10))

  def forward(self,xb):
    return xb @ weights + bias

model = Mnist_Logistic()
print(loss_func(model(xb), yb))

"""
Now it is easier to update model parameters:
"""
lr = 0.01
with torch.no_grad():
  for p in model.parameters(): p -= p.grad * lr
  model.zero_grad()


def fit():
  """
  Wrap our training loop in a fit function so we can run it again later
  """
  for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
      """
      Getting Batch
      """
      start_i = i * bs
      end_i = start_i + bs
      xb = x_train[start_i:end_i]
      yb = y_train[start_i:end_i]
      """ Prediction """
      pred = model(xb)
      """ Loss """
      loss = loss_func(pred, yb)
      """ Calculate gradients """
      loss.backward()
      """ Update parameters with gradients """
      with torch.no_grad():
        for p in model.parameters():
          p -= p.grad * lr
          model.zero_grad
print(loss_func(model(xb), yb))