Build Your First Neural Network with PyTorch

Site Admin · 11 Sep 2026 · 9 views

Build Your First Neural Network with PyTorch

PyTorch is a tensor library with autograd built in, so building a network is mostly declaring layers and letting the framework maintain gradients. This tutorial builds a small classifier that could handle MNIST digits after a few hundred lines of boring data plumbing.

Modeling a Network

A layer is a linear transform followed by an activation. In code, you define modules, stack them, and call the model like a function.

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = Net()
print(model)

nn.Linear stores its weights and biases; the forward method defines how data flows through them. Calling model(x) runs the forward pass.

Free Training Scaffolding

PyTorch ships components that remove most boilerplate. The flow is always the same:

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

for inputs, labels in loader:
    optimizer.zero_grad()
    logits = model(inputs)
    loss = criterion(logits, labels)
    loss.backward()
    optimizer.step()

zero_grad clears last step's gradients, the forward pass gives logits, the loss compares with labels, backward computes gradients, and step applies them. Five lines, repeated over every batch, is the entire training loop.

Zero to Sixty

Define the model, pick a loss, pick an optimizer, loop the data. Everything else, such as batches, shuffling, and metrics, is ceremony around these four decisions.

Move to GPU When Ready

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
for inputs, labels in loader:
    inputs, labels = inputs.to(device), labels.to(device)

Same code, different hardware, often a hundred times faster.

Key Points

  • nn.Module subclasses declare layers and forward.
  • Loss, optimizer, and a batch loop train anything.
  • Zero_grad, forward, backward, step is the rhythm.
  • Move tensors and model to the GPU together.
Share this post:

Comments (0)

Please login or register to comment.