Click here to download the full example code
Defining a Neural Network in PyTorch¶
Deep learning uses artificial neural networks (models), which are
computing systems that are composed of many layers of interconnected
units. By passing data through these interconnected units, a neural
network is able to learn how to approximate the computations required to
transform inputs into outputs. In PyTorch, neural networks can be
constructed using the
PyTorch provides the elegantly designed modules and classes, including
torch.nn, to help you create and train neural networks. An
nn.Module contains layers, and a method
In this recipe, we will use
torch.nn to define a neural network
intended for the MNIST
Before we begin, we need to install
torch if it isn’t already
pip install torch
Import all necessary libraries for loading our data
Define and initialize the neural network
Specify how data will pass through your model
[Optional] Pass data through your model to test
1. Import necessary libraries for loading our data¶
For this recipe, we will use
torch and its subsidiaries
import torch import torch.nn as nn import torch.nn.functional as F
2. Define and initialize the neural network¶
Our network will recognize images. We will use a process built into PyTorch called convolution. Convolution adds each element of an image to its local neighbors, weighted by a kernel, or a small matrix, that helps us extract certain features (like edge detection, sharpness, blurriness, etc.) from the input image.
There are two requirements for defining the
Net class of your model.
The first is writing an __init__ function that references
nn.Module. This function is where you define the fully connected
layers in your neural network.
Using convolution, we will define our model to take 1 input image channel, and output match our target of 10 labels representing numbers 0 through 9. This algorithm is yours to create, we will follow a standard MNIST algorithm.
class Net(nn.Module): def __init__(self): super(Net, self).__init__() # First 2D convolutional layer, taking in 1 input channel (image), # outputting 32 convolutional features, with a square kernel size of 3 self.conv1 = nn.Conv2d(1, 32, 3, 1) # Second 2D convolutional layer, taking in the 32 input layers, # outputting 64 convolutional features, with a square kernel size of 3 self.conv2 = nn.Conv2d(32, 64, 3, 1) # Designed to ensure that adjacent pixels are either all 0s or all active # with an input probability self.dropout1 = nn.Dropout2d(0.25) self.dropout2 = nn.Dropout2d(0.5) # First fully connected layer self.fc1 = nn.Linear(9216, 128) # Second fully connected layer that outputs our 10 labels self.fc2 = nn.Linear(128, 10) my_nn = Net() print(my_nn)
We have finished defining our neural network, now we have to define how our data will pass through it.
3. Specify how data will pass through your model¶
When you use PyTorch to build a model, you just have to define the
forward function, that will pass the data into the computation graph
(i.e. our neural network). This will represent our feed-forward
You can use any of the Tensor operations in the
class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 32, 3, 1) self.conv2 = nn.Conv2d(32, 64, 3, 1) self.dropout1 = nn.Dropout2d(0.25) self.dropout2 = nn.Dropout2d(0.5) self.fc1 = nn.Linear(9216, 128) self.fc2 = nn.Linear(128, 10) # x represents our data def forward(self, x): # Pass data through conv1 x = self.conv1(x) # Use the rectified-linear activation function over x x = F.relu(x) x = self.conv2(x) x = F.relu(x) # Run max pooling over x x = F.max_pool2d(x, 2) # Pass data through dropout1 x = self.dropout1(x) # Flatten x with start_dim=1 x = torch.flatten(x, 1) # Pass data through ``fc1`` x = self.fc1(x) x = F.relu(x) x = self.dropout2(x) x = self.fc2(x) # Apply softmax to x output = F.log_softmax(x, dim=1) return output
4. [Optional] Pass data through your model to test¶
To ensure we receive our desired output, let’s test our model by passing some random data through it.
# Equates to one random 28x28 image random_data = torch.rand((1, 1, 28, 28)) my_nn = Net() result = my_nn(random_data) print (result)
Each number in this resulting tensor equates to the prediction of the label the random tensor is associated to.
Congratulations! You have successfully defined a neural network in PyTorch.
Take a look at these other recipes to continue your learning:
Total running time of the script: ( 0 minutes 0.000 seconds)