You can construct neural networks using the torch.nn package. The torch.nn depends on autograd to define models and differentiate them. An nn.Module contains layers and a method forward(input) that returns the output.

The training process of a neural network:

  1. Define the neural network
  2. Iterate over a dataset (Inputs)
  3. Parse the input through the network
  4. Compute the loss function
  5. Perform backpropagation
  6. Update the weights of the network
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # 1 input channel, 6 output channels, 3x3 conv kernel
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)

        self.fc1 = nn.Linear(16 * 6 * 6, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # You only need to define the forward function, autograd automatically compute the backward function
        # Max pooling using a 2x2 window
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        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(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        
        return num_features

net = Net()
print(net)

params = list(net.parameters())
print(len(params))
print(params[0].size()) # Conv1's weight

# Input dat to the network and compute MSE loss with SGD optimiser
input = torch.randn(1, 1, 32, 32)
optimizer = optim.SGD(net.parameters(), lr = 0.01)
optimizer.zero_grad()
criterion = nn.MSELoss()
out = net(input)
target = torch.randn(10)
target = target.view(1, -1)

loss = criterion(out, target)
print(loss)
print(loss.grad_fn)
# print(loss.grad_fn.next_functions[0][0])

# Backpropagation
print('conv1.bias.grad before backward')
print(net.conv1.bias.grad)
loss.backward()
print('conv1.bias.grad after backward')
print(net.conv1.bias.grad)

# Update the weights
optimizer.step()
Ryan

Ryan

Data Scientist

Leave a Reply