Introduction to PyTorch

What Is PyTorch?

PyTorch is an open-source deep learning framework developed by Meta AI. It’s the most widely used framework in machine learning research and increasingly the standard for production systems. PyTorch’s key insight is treating neural network computations as standard Python code — tensors are just arrays, operations are just function calls, and debugging works with standard Python tools.

Installation

# CPU only
pip install torch torchvision

# With CUDA support (check pytorch.org for your CUDA version)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

Tensors — The Core Data Structure

import torch
import numpy as np

# Create tensors directly
x = torch.tensor([1.0, 2.0, 3.0, 4.0])
print(x.dtype)   # torch.float32
print(x.shape)   # torch.Size([4])
print(x.device)  # cpu

# From Python list
matrix = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
print(matrix.shape)   # torch.Size([2, 3])

# Factory functions — same as NumPy equivalents
zeros = torch.zeros(3, 4)
ones  = torch.ones(2, 3, dtype=torch.float32)
rand  = torch.rand(3, 3)                     # uniform [0, 1)
randn = torch.randn(3, 3)                    # standard normal
arange = torch.arange(0, 10, 2)             # [0, 2, 4, 6, 8]
linspace = torch.linspace(0, 1, 5)          # [0.0, 0.25, 0.5, 0.75, 1.0]

# From NumPy — shares memory (no copy)
np_arr = np.array([1., 2., 3.])
t = torch.from_numpy(np_arr)   # shares memory
t2 = torch.tensor(np_arr)     # copies memory

# To NumPy
arr = t.numpy()

Tensor Operations

import torch

a = torch.tensor([[1., 2.], [3., 4.]])
b = torch.tensor([[5., 6.], [7., 8.]])

# Arithmetic — all element-wise
print(a + b)
print(a * b)
print(a ** 2)
print(torch.sqrt(a))

# Matrix multiplication
print(a @ b)              # same as torch.matmul(a, b)
print(torch.mm(a, b))     # 2-D only

# Aggregations
print(a.sum())            # scalar
print(a.sum(dim=0))       # sum along rows → shape (2,)
print(a.mean(dim=1))      # mean along columns → shape (2,)
print(a.max())
print(a.argmax())

# Reshape
x = torch.arange(12.)
print(x.reshape(3, 4))
print(x.view(3, 4))       # view shares memory — only works on contiguous tensors
print(x.reshape(3, -1))   # -1 infers the missing dimension

# Indexing — same as NumPy
print(a[0])         # first row
print(a[:, 1])      # second column
print(a[a > 2])     # boolean mask

# Type casting
x = torch.tensor([1, 2, 3])
print(x.dtype)              # torch.int64
print(x.float())            # convert to float32
print(x.to(torch.float32))  # explicit cast

Moving to GPU

import torch

# Check GPU availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Move tensor to GPU
x = torch.randn(1000, 1000)
x_gpu = x.to(device)          # or x.cuda()

# Operations run on the device where the tensor lives
y_gpu = x_gpu @ x_gpu.T       # matrix multiply on GPU

# Move back to CPU for NumPy operations
y_cpu = y_gpu.cpu()
arr = y_cpu.numpy()

# Device-aware code pattern
x = torch.randn(100, 10).to(device)
w = torch.randn(10, 5).to(device)
result = x @ w   # both on same device — works

Autograd — Automatic Differentiation

Autograd is PyTorch’s automatic differentiation engine. It records operations on tensors that have requires_grad=True, then computes gradients by traversing the computation graph in reverse (backpropagation).

import torch

# requires_grad=True: track all operations on this tensor
x = torch.tensor(3.0, requires_grad=True)

# Build a computation: y = x^2 + 2x + 1
y = x ** 2 + 2 * x + 1

# Compute gradients — dy/dx = 2x + 2 = 2(3) + 2 = 8
y.backward()   # backpropagate
print(x.grad)  # tensor(8.)

# Gradient accumulates — zero before next forward pass
x.grad.zero_()  # in-place zero (convention for in-place ops: trailing _)

import torch

# Gradient of a vector function
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = (x ** 2).sum()   # y = x1^2 + x2^2 + x3^2

y.backward()
print(x.grad)   # dy/dx_i = 2*x_i → [2., 4., 6.]

# Detach from computation graph (stop tracking gradients)
z = x.detach()   # z shares storage but has no gradient history

# Context manager to temporarily disable gradient tracking
with torch.no_grad():
    # Runs faster, no memory for graph — use for inference
    result = x ** 2

Your First Neural Network with autograd

import torch
import torch.nn.functional as F

# A manual 2-layer neural network — no nn.Module, just tensors and autograd
torch.manual_seed(42)

# XOR dataset
X = torch.tensor([[0., 0.], [0., 1.], [1., 0.], [1., 1.]])
y = torch.tensor([[0.], [1.], [1.], [0.]])

# Weights — initialized randomly, requires_grad to learn
W1 = torch.randn(2, 4, requires_grad=True) * 0.3
b1 = torch.zeros(4, requires_grad=True)
W2 = torch.randn(4, 1, requires_grad=True) * 0.3
b2 = torch.zeros(1, requires_grad=True)

lr = 0.1

for epoch in range(1000):
    # Forward pass
    hidden = F.relu(X @ W1 + b1)       # (4, 4)
    output = torch.sigmoid(hidden @ W2 + b2)  # (4, 1)

    # Loss: binary cross-entropy
    loss = F.binary_cross_entropy(output, y)

    # Backward pass — compute gradients for all requires_grad tensors
    loss.backward()

    # Gradient descent step (no_grad: parameter update shouldn't be tracked)
    with torch.no_grad():
        W1 -= lr * W1.grad
        b1 -= lr * b1.grad
        W2 -= lr * W2.grad
        b2 -= lr * b2.grad

    # Zero gradients before the next forward pass
    W1.grad.zero_()
    b1.grad.zero_()
    W2.grad.zero_()
    b2.grad.zero_()

    if (epoch + 1) % 200 == 0:
        print(f"Epoch {epoch+1:4d} | Loss: {loss.item():.4f}")

# Final predictions
with torch.no_grad():
    preds = (torch.sigmoid(F.relu(X @ W1 + b1) @ W2 + b2) > 0.5).float()
    print(f"Predictions: {preds.flatten().tolist()}")  # [0, 1, 1, 0]