Build Your First Neural Network in PyTorch (2026)

 


What Is PyTorch? A Complete Beginner's Guide to Deep Learning in Python (2026)

Part 3: Building Your First Neural Network with PyTorch

Introduction

In Part 1, you learned what PyTorch is, its history, why it's popular, and how to install it.

In Part 2, you explored tensors, tensor operations, CPU vs GPU, and PyTorch's automatic differentiation (Autograd).

Now it's time to put everything together and build your first neural network.

A neural network is the heart of modern AI systems. Whether you're classifying images, translating languages, generating text, or predicting stock prices, neural networks make these intelligent applications possible.

PyTorch makes building neural networks remarkably simple through the powerful torch.nn module. In this chapter, you'll learn how neural networks are structured, how to define your own model, how the forward() function works, how loss functions and optimizers train the model, how the training loop operates, and how to save and reload trained models.

Let's dive in.

Building a Neural Network

A neural network is a collection of connected layers that learn patterns from data.

A simple neural network consists of:

  • Input Layer – Receives the data.
  • Hidden Layer(s) – Learns patterns and relationships.
  • Output Layer – Produces the prediction.

For example, imagine a model that predicts whether an email is spam.

Email Features
       │
       ▼
Input Layer
       │
       ▼
Hidden Layers
       │
       ▼
Output Layer
    Spam / Not Spam

During training, the network gradually adjusts its internal parameters (called weights and biases) so that its predictions become more accurate.

The torch.nn Module

PyTorch provides the torch.nn module, which contains pre-built components for creating neural networks.

It includes:

  • Linear (fully connected) layers
  • Convolutional layers
  • Recurrent layers
  • Activation functions
  • Dropout layers
  • Batch normalization
  • Loss functions
  • Containers for building models

Instead of implementing these from scratch, you simply import the module:

import torch
import torch.nn as nn

The nn alias is widely used because it keeps code concise and readable.

Understanding nn.Module

Every custom neural network in PyTorch inherits from nn.Module.

Think of nn.Module as the blueprint that provides everything a neural network needs:

  • Stores trainable parameters
  • Tracks layers
  • Supports saving/loading
  • Moves models between CPU and GPU
  • Switches between training and evaluation modes

Creating a Simple Neural Network

import torch
import torch.nn as nn

class SimpleNet(nn.Module):

    def __init__(self):
        super().__init__()

        self.layer1 = nn.Linear(2, 4)
        self.layer2 = nn.Linear(4, 1)

    def forward(self, x):
        x = self.layer1(x)
        x = torch.relu(x)
        x = self.layer2(x)

        return x

Although this code may seem unfamiliar at first, each part has a specific purpose.

Breaking Down the Model

The Constructor (__init__)

def __init__(self):
    super().__init__()

This initializes the parent nn.Module.

Without this line, PyTorch wouldn't know how to manage your neural network.

Defining Layers

self.layer1 = nn.Linear(2,4)

This creates a fully connected layer.

It accepts:

  • 2 input values
  • Produces 4 output values

The second layer:

self.layer2 = nn.Linear(4,1)

accepts 4 inputs and produces one prediction.

Understanding the forward() Function

The forward() method defines how information flows through the neural network.

def forward(self, x):

    x = self.layer1(x)

    x = torch.relu(x)

    x = self.layer2(x)

    return x

When you call:

output = model(input)

PyTorch automatically executes the forward() method.

You never call forward() directly.

Activation Functions

Without activation functions, a neural network would behave like a simple linear equation.

One of the most common activation functions is ReLU.

torch.relu(x)

ReLU stands for:

Rectified Linear Unit

Its rule is simple:

If x > 0
    output = x

Else
    output = 0

ReLU helps neural networks learn complex, non-linear patterns.

Other popular activation functions include:

  • Sigmoid
  • Tanh
  • Softmax
  • Leaky ReLU
  • GELU

Creating the Model

Once the class is defined:

model = SimpleNet()

print(model)

Output:

SimpleNet(
  (layer1): Linear(in_features=2, out_features=4)
  (layer2): Linear(in_features=4, out_features=1)
)

PyTorch neatly displays the model architecture.

Making Predictions

Let's pass some input through the model.

import torch

sample = torch.tensor([[2.0,3.0]])

prediction = model(sample)

print(prediction)

Initially, the output will appear random.

That's because the network hasn't learned anything yet.

Training fixes that.

What Is a Loss Function?

A neural network needs a way to measure how wrong its predictions are.

This measurement is called the loss.

A loss function compares:

  • Predicted output
  • Correct answer

The smaller the loss, the better the model.

Mean Squared Error (MSE)

Used for regression problems.

criterion = nn.MSELoss()

Formula:

Loss = Average of (Prediction − Target)²

If prediction:

8

Actual answer:

10

Loss:

(8−10)² = 4

Training attempts to reduce this value.

Cross Entropy Loss

Used for classification.

criterion = nn.CrossEntropyLoss()

Applications:

  • Cat vs Dog
  • Spam Detection
  • Sentiment Analysis
  • Digit Recognition

It is one of the most commonly used loss functions in deep learning.

Optimizers

Knowing the error isn't enough.

The model must also update its parameters to reduce that error.

This is the optimizer's job.

An optimizer changes the model's weights after every training step.

Stochastic Gradient Descent (SGD)

One of the oldest optimizers.

optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.01
)

Here:

lr

means Learning Rate.

A learning rate of:

0.01

means the optimizer makes relatively small adjustments after each step.

Adam Optimizer

Adam is one of the most popular optimizers because it generally converges faster and requires less manual tuning.

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=0.001
)

Many beginners start with Adam because it performs well across a wide range of problems.

The Training Loop

Training follows the same basic sequence repeatedly.

  1. Make predictions.
  2. Calculate the loss.
  3. Compute gradients.
  4. Update the weights.
  5. Repeat.

In PyTorch, this is implemented with a training loop.

Example Training Loop

for epoch in range(100):

    predictions = model(inputs)

    loss = criterion(predictions, targets)

    optimizer.zero_grad()

    loss.backward()

    optimizer.step()

Let's understand each step.

Step 1: Forward Pass

predictions = model(inputs)

The input data passes through the neural network, producing predictions.

Step 2: Calculate Loss

loss = criterion(predictions, targets)

The loss function measures how far the predictions are from the correct answers.

Step 3: Clear Old Gradients

optimizer.zero_grad()

PyTorch accumulates gradients by default. Calling zero_grad() clears the gradients from the previous iteration.

Step 4: Backpropagation

loss.backward()

Autograd computes the gradients of the loss with respect to every trainable parameter in the model.

Step 5: Update Parameters

optimizer.step()

The optimizer uses those gradients to update the model's weights and biases.

After many iterations, the model gradually becomes better at making predictions.

Full Example

import torch
import torch.nn as nn

class SimpleNet(nn.Module):

    def __init__(self):
        super().__init__()

        self.fc = nn.Linear(2,1)

    def forward(self,x):
        return self.fc(x)

model = SimpleNet()

criterion = nn.MSELoss()

optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.01
)

inputs = torch.tensor([[1.,2.],[2.,3.]])
targets = torch.tensor([[3.],[5.]])

for epoch in range(200):

    outputs = model(inputs)

    loss = criterion(outputs,targets)

    optimizer.zero_grad()

    loss.backward()

    optimizer.step()

print("Training Complete")

Although this is a tiny example, it demonstrates the complete workflow used in much larger deep learning projects.

Saving a Trained Model

Training can take minutes, hours, or even days. Instead of retraining every time, you can save the learned parameters.

torch.save(
    model.state_dict(),
    "model.pth"
)

The file model.pth stores the model's learned weights.

Loading a Saved Model

To use the model later:

model = SimpleNet()

model.load_state_dict(
    torch.load("model.pth")
)

model.eval()

Calling model.eval() switches the network to evaluation mode, ensuring layers such as Dropout and Batch Normalization behave correctly during inference.

Best Practices

When building PyTorch models:

  • Use meaningful variable names.
  • Start with simple architectures before adding complexity.
  • Split data into training and validation sets.
  • Monitor training loss to detect problems.
  • Save checkpoints regularly.
  • Use GPU acceleration when available.
  • Keep your code modular and well documented.

Frequently Asked Questions (FAQs)

Is PyTorch free?

Yes. PyTorch is open-source and free for both personal and commercial use.

Is PyTorch beginner-friendly?

Yes. Its Pythonic syntax and extensive documentation make it one of the easiest deep learning frameworks to learn.

Does PyTorch support GPUs?

Yes. PyTorch supports CUDA-enabled NVIDIA GPUs and can also run efficiently on CPUs.

What's the difference between PyTorch and TensorFlow?

Both are powerful deep learning frameworks. PyTorch is often preferred for research and experimentation because of its flexibility and intuitive programming model, while TensorFlow has long been popular for large-scale production deployments. Today, both frameworks are widely used in research and industry.

Do I need advanced mathematics to start?

No. A basic understanding of algebra and functions is enough to begin. As you progress, learning topics such as linear algebra, probability, and calculus will help you understand deep learning more deeply.

Conclusion

Congratulations! You've completed this three-part beginner's guide to PyTorch.

You started by learning what PyTorch is and why it's become one of the most widely used deep learning frameworks. You explored tensors, automatic differentiation, CPU vs GPU computation, and finally built your first neural network using nn.Module.

While the examples in this guide are intentionally simple, they introduce the same core concepts used in advanced AI systems such as image classifiers, recommendation engines, large language models, and computer vision applications.

The best way to continue learning is by building small projects, experimenting with different datasets, and gradually exploring more advanced topics like convolutional neural networks (CNNs), recurrent neural networks (RNNs), transformers, transfer learning, and model deployment. With consistent practice, you'll develop the skills needed to create powerful AI applications using PyTorch.

learn from step-1:
https://khayyamshah2007.blogspot.com/2026/08/what-is-pytorch-complete-beginners.html

Comments

Post a Comment