PyTorch is an open-source machine learning library used for a wide variety of tasks such as deep learning, natural language processing (NLP), and computer vision. It provides a flexible platform to build machine learning models and comes with strong support for GPU acceleration, making it popular among researchers and developers.
torchvision and torchtext libraries for various computer vision and NLP tasks.How to use PyTorch to create a Neural Network for classifying MNIST digits:
#!/usr/bin/env python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
# Define a simple neural network
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(28*28, 128) # Fully connected layer 1
self.fc2 = nn.Linear(128, 10) # Fully connected layer 2 (10 classes)
def forward(self, x):
x = x.view(-1, 28*28) # Flatten the image
x = torch.relu(self.fc1(x)) # Apply ReLU to fc1
x = self.fc2(x) # Final output (logits)
return x
# Load dataset and preprocess
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
# Initialize the model, loss function, and optimizer
net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)
# Training loop
for epoch in range(2):
running_loss = 0.0
for inputs, labels in trainloader:
optimizer.zero_grad() # Zero the gradients
outputs = net(inputs) # Forward pass
loss = criterion(outputs, labels) # Compute the loss
loss.backward() # Backward pass
optimizer.step() # Optimize
running_loss += loss.item()
print(f'Epoch {epoch+1}, Loss: {running_loss / len(trainloader)}')
print("Finished Training")
torchvision: Contains datasets, models, and transforms for computer vision.torchtext : For NLP models and datasets.torchaudio : For Audio and Speech Processing.PyTorch is a Python-first deep-learning framework originally developed at Meta AI and now governed by the PyTorch Foundation. Its core object is the tensor, an n-dimensional array that can live on CPU or GPU and that records operations for automatic differentiation. On top of tensors, torch.nn provides layers and loss functions, torch.optim provides optimisers, and torch.utils.data provides datasets and batched loaders. Because models are ordinary Python classes executed eagerly, debugging uses normal tools such as print statements and pdb, which is a large part of why PyTorch became the dominant framework in research and is widely used in production through TorchScript, torch.compile, ONNX export and serving stacks.
Every PyTorch training script follows the same loop, visible in the MNIST example above. A DataLoader yields mini-batches. The model's forward method turns inputs into outputs, and a loss function reduces them against labels to a single number. Calling backward() on that number computes the gradient of the loss with respect to every parameter. The optimiser then uses those gradients to update the parameters, and the gradients are zeroed before the next batch. Evaluation wraps the forward pass in torch.no_grad() and switches the model to eval() mode so that dropout and batch normalisation behave deterministically.
When a tensor has requires_grad=True, each operation on it creates a node with a grad_fn that knows how to compute its local derivative. The nodes form a directed acyclic graph from the leaf parameters to the loss. backward() applies the chain rule through that graph and stores results in each leaf's .grad. The graph is freed after the backward pass and rebuilt on the next forward pass. This define-by-run approach is what "dynamic computational graph" means: variable-length sequences, recursion and data-dependent branching need no special graph operators.
Around the core library sits a wide ecosystem. torchvision, torchaudio and the Hugging Face transformers library provide pretrained models and data pipelines; torch.distributed and Fully Sharded Data Parallel scale training across GPUs and nodes; and torch.compile, introduced in PyTorch 2.0, captures Python code into optimised graphs and kernels while keeping the eager programming model. For deployment, models can be exported with torch.export or to ONNX and served with tools such as TorchServe, Triton or ONNX Runtime.
import torch
import torch.nn as nn
device = "cuda" if torch.cuda.is_available() else "cpu"
# autograd on a scalar function: d/dw (w*3 + 1)^2 at w=2 -> 2*(7)*3 = 42
w = torch.tensor(2.0, requires_grad=True)
y = (w * 3 + 1) ** 2
y.backward()
print(w.grad) # tensor(42.)
# fit y = 2x + 1 with a single Linear layer
x = torch.linspace(-1, 1, 200, device=device).unsqueeze(1)
t = 2 * x + 1 + 0.05 * torch.randn_like(x)
model = nn.Linear(1, 1).to(device)
opt = torch.optim.Adam(model.parameters(), lr=0.1)
for step in range(300):
opt.zero_grad()
loss = nn.functional.mse_loss(model(x), t)
loss.backward()
opt.step()
print(model.weight.item(), model.bias.item()) # close to 2.0 and 1.0
# save and reload weights
torch.save(model.state_dict(), "linear.pt")
m2 = nn.Linear(1, 1); m2.load_state_dict(torch.load("linear.pt"))
The first block shows autograd computing an exact derivative. The second is the full training loop in miniature, running on GPU if one is present. Saving the state_dict rather than the whole model object is the recommended practice because it does not depend on pickling class definitions.
| Framework | Style | Strengths | Typical use |
|---|---|---|---|
| PyTorch | Eager, define-by-run | Debuggability, research ecosystem, Hugging Face default | Research, LLMs, most new projects |
| TensorFlow / Keras | Keras high-level API, graph via tf.function | Mature deployment tools (TF Lite, TF Serving) | Existing TF stacks, mobile |
| JAX | Functional, jit-compiled | Composable transforms, TPU performance | Large-scale research on TPUs |
| scikit-learn | Classical ML estimators | Simple API, no GPU needed | Tabular data, baselines |
optimizer.zero_grad(), so gradients accumulate across batches.loss instead of loss.item() in a running total, which keeps every graph alive and leaks memory.x.view on non-contiguous tensors; use reshape.