Neural Networks PyTorch
Tensor Autograd

PyTorch for Neural Networks

PyTorch centers on tensors (multi-dimensional arrays) that can live on CPU or GPU. Autograd tracks operations on tensors that require gradients, building a graph for reverse-mode differentiation. You define models as subclasses of nn.Module with learnable parameters in nn.Parameter or child modules; forward() describes the computation. A typical training step: loss.backward(), optimizer.step(), optimizer.zero_grad().

nn.Module DataLoader .to(device) torch.save

Tensors & Device

Create tensors with torch.tensor, torch.randn, etc. Move data and model with .to(device) where device = torch.device("cuda" if torch.cuda.is_available() else "cpu"). Use model.train() / model.eval() to toggle dropout and batch norm behavior.

nn.Module & Training Loop

Minimal training loop 
import torch
import torch.nn as nn

model = MyModel().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

for xb, yb in loader:
    xb, yb = xb.to(device), yb.to(device)
    logits = model(xb)
    loss = loss_fn(logits, yb)
    opt.zero_grad()
    loss.backward()
    opt.step()
Wrap training iteration in torch.cuda.amp.autocast() and GradScaler when using mixed precision on NVIDIA GPUs.

Dataset & DataLoader

Subclass torch.utils.data.Dataset to implement __len__ and __getitem__, or use ImageFolder, TensorDataset. DataLoader batches, shuffles, and uses multiple workers for parallel loading.

Summary

  • PyTorch = tensors + autograd + nn modules + optimizers.
  • Keep device placement consistent; use train/eval modes correctly.
  • DataLoader feeds the loop; checkpoint with torch.save / load_state_dict.
  • Next: TensorFlow / Keras for the other major ecosystem.

Prefer a high-level API? Keras inside TensorFlow offers model.fit and saved models.

Previous: Metrics Next: TensorFlow / Keras