deep learning practice

import torch
import torch.nn as nn
import torchvision

# Start coding here... 

class AlexNet(nn.Module):
  def __init__(self, num_classes=1000):
      super(AlexNet, self).__init__()
      self.conv1 = nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2)
      self.relu = nn.ReLU(inplace=True)
      self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2)
      self.conv2 = nn.Conv2d(64, 192, kernel_size=5, padding=2)
      self.conv3 = nn.Conv2d(192, 384, kernel_size=3, padding=1)
      self.conv4 = nn.Conv2d(384, 256, kernel_size=3, padding=1)
      self.conv5 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
      self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
      self.fc1 = nn.Linear(256 * 6 * 6, 4096)
      self.fc2 = nn.Linear(4096, 4096)
      self.fc3 = nn.Linear(4096, num_classes)

  def forward(self, x):
      x = self.relu(self.conv1(x))
      x = self.maxpool(x)
      x = self.relu(self.conv2(x))
      x = self.maxpool(x)
      x = self.relu(self.conv3(x))
      x = self.relu(self.conv4(x))
      x = self.relu(self.conv5(x))
      x = self.maxpool(x)
      x = self.avgpool(x)
      x = x.view(x.size(0), 256 * 6 * 6)
      x = self.relu(self.fc1(x))
      x = self.relu(self.fc2(x))
      return self.fc3(x)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        
        # Instantiate two convolutional layers
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=5, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(in_channels=5, out_channels=10, kernel_size=3, padding=1)
        
        # Instantiate the ReLU nonlinearity
        self.relu = nn.ReLU(inplace=True)
        
        # Instantiate a max pooling layer
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # Instantiate a fully connected layer
        self.fc = nn.Linear(49 * 10, 10)
        
    def forward(self, x):
        # Apply conv followed by relu, then in next line pool
        x = self.relu(self.conv1(x))
        x = self.pool(x)
        
        # Apply conv followed by relu, then in next line pool
        x = self.relu(self.conv2(x))
        x = self.pool(x)
        
        # Prepare the image for the fully connected layer
        x = x.view(-1, 7 * 7 * 10)
        
        # Apply the fully connected layer and return the result
        return self.fc(x)

import torchvision.transforms as transforms
from torchvision.datasets import MNIST

# Transform the data to torch tensors and normalize it
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307), (0.3081))
])

# Preparing the training and test set
trainset = torchvision.datasets.MNIST('mnist', train=True, transform=transform)
testset = torchvision.datasets.MNIST('mnist', train=True, transform=transform)

# Prepare loader
train_loader = torch.utils.data.DataLoader(trainset, batch_size=1, shuffle=True, num_workers=0)
test_loader = torch.utils.data.DataLoader(testset, batch_size=1, shuffle=False, num_workers=0)

import torch.optim as optim

net = Net()
optimizer = optim.Adam(net.parameters(), lr=3e-4)
criterion = nn.CrossEntropyLoss()

for i, data in enumerate(train_loader, 0):
    inputs, labels = data
    optimizer.zero_grad()
    
    # Compute the forward pass
    outputs = net(inputs)
    
    # Compute the loss function
    loss = criterion(outputs, labels)
    
    # Compute the gradients
    loss.backward()
    
    # Update the weights
    optimizer.step()

net.eval()
# Iterate over the data in the test_loader
for i, data in enumerate(test_loader):
    # Get the image and label from data
    image, label = data
    
    # Make a forward pass in the net with your image
    output = net(image)
    
    # Argmax the results of the net
    _, predicted = torch.max(output.data, 1)
    
    if predicted == label:
        print("Yipes, your net made the right prediction " + str(predicted))
    else:
        print("Your net prediction was " + str(predicted) + ", but the correct label is: " + str(label))
        
    if i > 10:
        break