PyTorch is one of the most used libraries for building deep learning models, especially neural network-based models. In many tasks related to deep learning, we find the use of PyTorch because of its features and capabilities like production-ready, distributed training, robust ecosystem, and cloud support. In this article, we will learn how we can build a simple neural network using the PyTorch library in just a few steps. For this purpose, we will demonstrate a hands-on implementation where we will build a simple neural network for a classification problem.
We will accomplish this implementation in the following steps:-
- Step 1: Creating the data
- Step 2: Loading the data using data loader
- Step 3: Building a neural network model
- Defining the neural net class
- Instantiating the classifier
- Step 4: Training the neural network model
- Optimizing loss curve
- Defining decision boundaries
- Step 5: Making Predictions
Let’s start with the first step, where we will create a dataset for implementation.
AIM Daily XO
Join our editors every weekday evening as they steer you through the most significant news of the day, introduce you to fresh perspectives, and provide unexpected moments of joy
Your newsletter subscriptions are subject to AIM Privacy Policy and Terms and Conditions.
Step 1: Creating the data
In this article, we will understand the intuition behind defining and working the neural networks in PyTorch. For this purpose, we will build a neural network from scratch and apply it to a classification problem. For this purpose, we will create a dataset using make_classification under sklearn.datasets which we use in making a classification data set.
Using the following lines of code we can create the dataset for classification:
Download our Mobile App
from sklearn.datasets import make_classification
X, Y = make_classification(
n_features=4, n_redundant=0, n_informative=3, n_clusters_per_class=2, n_classes=3
)
Using the above code we have created a dataset for classification wherein the dataset we have 4 features with 3 informative features and 3 classes.
Let’s visualize the dataset.
import matplotlib.pyplot as plt
plt.title("Multi-class data, 4 informative features, 3 classes", fontsize="large")
plt.scatter(X[:, 0], X[:, 1], marker="o", c=Y, s=25, edgecolor="k")

Here we can see our dataset in a two-dimensional space and points are clearly visible from three classes.
After making the dataset, we are ready to build a classification model. Since in this article, we are discussing a simple implementation of a neural network using the PyTorch, we will use a two-layer neural network where we can use sigmoid as our activation function. Data points in the above graph will be our input coordinates and classes related to the dots are the ground truth.
Step 2: Loading the data using data loader
Before defining the model, we are required to split our dataset into tests and train sets. That can be done using the following lines of codes.
from sklearn.model_selection import train_test_split
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.33, random_state=42)
Let’s check the shape of the split dataset.
X_train.shape, X_test.shape, Y_train.shape, Y_test.shape
Output:

After this, we are required to define the train data as a PyTorch tensor.
import torch
Y_test = torch.from_numpy(X_test)
Y_test = torch.from_numpy(np.asarray(Y_test))
To use the batches of the dataset we are required to put the data through the Dataloader module of PyTorch. Using the below lines of code we can do that.
from torch.utils.data import Dataset, DataLoader
class Data(Dataset):
def __init__(self):
self.X=torch.from_numpy(X_train)
self.Y=torch.from_numpy(Y_train)
self.len=self.X.shape[0]
def __getitem__(self,index):
return self.X[index], self.Y[index]
def __len__(self):
return self.len
data=Data()
loader=DataLoader(dataset=data,batch_size=64)
Checking the shape of the data under the tensor.
print(data.X[0:5])
print(data.X.shape)
print(data.Y[0:5])
print(data.Y.shape)
Output:

Step 3: Building a neural network model
First of all, we will define the dimensions of the network.
input_dim=4 # how many Variables are in the dataset
hidden_dim = 25 # hidden layers
output_dim=3 # number of classes
Step 3.1: Defining the neural net class
A class for creating simple networks can be defined using the following lines of codes.
import torch.nn as nn
class Net(nn.Module):
def __init__(self,input,H,output):
super(Net,self).__init__()
self.linear1=nn.Linear(input,H)
self.linear2=nn.Linear(H,output)
def forward(self,x):
x=torch.sigmoid(self.linear1(x))
x=self.linear2(x)
return x
Step 3.2: Instantiating the classifier
Now with the above-defined dimensions, we can instantiate a model instance using the class defined in the last step.
clf=Net(input_dim,hidden_dim,output_dim)
Let’s check the information about the model
print(clf.parameters)
Output:

Here we can see that the model which we have prepared will take 4 features as input and will give 4 features as output.
Step 4: Training the neural network model
Before training, we are required to define the criterion to calculate gradients of parameters and optimizers to update the parameters.
criterion=nn.CrossEntropyLoss()
optimizer=torch.optim.SGD(clf.parameters(), lr=0.1)
After defining the criterion and optimizer, we are ready to train our model. Using the following lines of codes we can train it.
learning_rate = 1e-1
loss_list = []
for t in range(1000):
y_pred = clf(x)
loss = criterion(y_pred, y)
loss_list.append(loss.item())
clf.zero_grad()
loss.backward()
with torch.no_grad():
for param in clf.parameters():
param -= learning_rate * param.grad
In this training, we have defined methods that will get the accumulated gradient to zero, append the loss on the loss_list, and will get a new gradient, and update parameters using the backward propagation.
Step 4.1: Optimizing loss curve
The loss curve will let us know about the performance of the model.
Visualizing loss curve
step = np.linspace(0,1000,1000)
plt.plot(step,np.array(loss_list))
Output:

In the above output, we can see that our loss curve converged, and using this model we can also define our decision boundaries.
Step 4.2: Defining decision boundaries
params = list(clf.parameters())
w = params[0].detach().numpy()[0]
b = params[1].detach().numpy()[0]
t= params[3].detach().numpy()[0]
plt.scatter(X[:, 0], X[:, 1], c=Y,cmap='jet')
u = np.linspace(X[:, 0].min(), X[:, 0].max(), 2)
plt.plot(u, (0.5-b-w[0]*u)/w[1])
plt.plot(u, (0.5-t-w[0]*u)/w[1])
plt.xlim(X[:, 0].min()-0.5, X[:, 0].max()+0.5)
plt.ylim(X[:, 1].min()-0.5, X[:, 1].max()+0.5)
Output:

In the above output, we can see that we have defined our decision boundaries and it is looking good. Now we can use the model for making predictions.
Making predictions
Let’s make the predictions with the trained neural network.
x_val = torch.from_numpy(X_test)
z=clf(x_val)
yhat=torch.max(z.data,1)
yhat[1]
Output:

In the above output, we can see the predictions of the class labels of the given input patterns. Using a few of the modifications like the number of layers or number of epochs we can improve the performance of our model.
Final words
In this article, we have learnt how to define a neural net in PyTorch in just a few steps. We have defined a dataset on which we have trained a double-layer neural network. We have used this network for classification and implemented that in PyTorch. This can be an example of a simple multiclass classification neural network also.