Introduction to PyTorch

1┌─────────────────────────┬─────────────────────────┐
2│       TENSORFLOW        │         PYTORCH         │
3├─────────────────────────┼─────────────────────────┤
4│ Static computation graph│ Dynamic computation graph│
5│ Define-and-run          │ Define-by-run           │
6│ More production tools   │ More research-friendly  │
7│ tf.data for data loading│ DataLoader for batching │
8│ Keras high-level API    │ nn.Module architecture  │
9│ TensorBoard built-in    │ TensorBoard supported   │
10└─────────────────────────┴─────────────────────────┘

1# TensorFlow (Static Graph - Eager mode changed this)
2# Historically required building graph first
3@tf.function
4def forward(x):
5    return x * 2  # Graph compiled once
6
7# PyTorch (Dynamic Graph)
8def forward(x):
9    if x.sum() > 0:  # Can use Python control flow!
10        return x * 2
11    else:
12        return x * 3  # Different path based on data

1┌─────────────────────────────────────────┐
2│           PYTORCH LAYERS                │
3├─────────────────────────────────────────┤
4│  torch.nn         │ Neural network API  │
5│  torch.optim      │ Optimizers (SGD,    │
6│                   │   Adam, etc.)       │
7├─────────────────────────────────────────┤
8│  torch.autograd   │ Automatic           │
9│                   │ differentiation     │
10├─────────────────────────────────────────┤
11│  torch.Tensor     │ Multi-dimensional   │
12│                   │ arrays with GPU     │
13├─────────────────────────────────────────┤
14│  C++ Backend      │ ATen (C++ tensor    │
15│  (libTorch)       │   library)          │
16└─────────────────────────────────────────┘

1import torch
2import torch.nn as nn
3
4print(f"PyTorch version: {torch.__version__}")
5print(f"CUDA available: {torch.cuda.is_available()}")
6
7# Create a tensor (PyTorch's fundamental data structure)
8x = torch.tensor([1.0, 2.0, 3.0])
9print(f"Tensor: {x}")
10
11# Simple neural network
12class SimpleNet(nn.Module):
13    def __init__(self):
14        super().__init__()
15        self.layer = nn.Linear(3, 1)
16
17    def forward(self, x):
18        return self.layer(x)
19
20# Create model
21model = SimpleNet()
22output = model(x)
23print(f"Model output: {output}")

1# PyTorch feels like NumPy with GPU acceleration
2x = torch.randn(100, 10)  # Like np.random.randn
3y = x.mean(dim=0)         # Like x.mean(axis=0)
4z = torch.softmax(y, dim=0)  # Built-in operations

1# Use regular Python debugger
2import pdb
3
4def forward(x):
5    pdb.set_trace()  # Drop into debugger
6    return x * 2
7
8# Or print anywhere
9x = torch.randn(5)
10print(f"Debug: {x}")  # Just works!

1# Easy to extend with custom operations
2class CustomLayer(nn.Module):
3    def forward(self, x):
4        # Any Python code!
5        if self.training:
6            return x * torch.randn_like(x)  # Custom behavior
7        return x

1import torch
2import numpy as np
3
4# From Python lists
5tensor1 = torch.tensor([1, 2, 3, 4])
6print(f"From list: {tensor1}")
7
8# From NumPy arrays (share memory!)
9numpy_array = np.array([1, 2, 3, 4])
10tensor2 = torch.from_numpy(numpy_array)
11print(f"From NumPy: {tensor2}")
12
13# Random initialization
14random_tensor = torch.randn(3, 4)  # Normal distribution
15uniform_tensor = torch.rand(3, 4)   # Uniform [0, 1)
16zeros = torch.zeros(2, 3)
17ones = torch.ones(2, 3)
18
19print(f"Random tensor shape: {random_tensor.shape}")
20
21# Like NumPy arrays
22similar = torch.randn_like(random_tensor)
23print(f"Similar tensor: {similar.shape}")
24
25# Specify data type
26float_tensor = torch.tensor([1, 2, 3], dtype=torch.float32)
27int_tensor = torch.tensor([1.5, 2.5, 3.5], dtype=torch.int64)
28print(f"Float tensor: {float_tensor}")
29print(f"Int tensor: {int_tensor}")

1x = torch.randn(3, 4, 5)
2
3print(f"Shape: {x.shape}")          # torch.Size([3, 4, 5])
4print(f"Size: {x.size()}")          # Same as shape
5print(f"Dtype: {x.dtype}")          # torch.float32
6print(f"Device: {x.device}")        # cpu or cuda
7print(f"Requires grad: {x.requires_grad}")  # False by default
8print(f"Total elements: {x.numel()}")  # 3*4*5 = 60

1# Element-wise operations
2a = torch.tensor([1.0, 2.0, 3.0])
3b = torch.tensor([4.0, 5.0, 6.0])
4
5print(f"Addition: {a + b}")
6print(f"Multiplication: {a * b}")
7print(f"Power: {a ** 2}")
8
9# Matrix operations
10matrix1 = torch.randn(3, 4)
11matrix2 = torch.randn(4, 5)
12
13# Matrix multiplication
14result = torch.mm(matrix1, matrix2)  # Or matrix1 @ matrix2
15print(f"Matrix multiply result shape: {result.shape}")  # [3, 5]
16
17# Batch matrix multiplication
18batch1 = torch.randn(10, 3, 4)  # 10 matrices of 3x4
19batch2 = torch.randn(10, 4, 5)  # 10 matrices of 4x5
20batch_result = torch.bmm(batch1, batch2)  # [10, 3, 5]
21
22# Reshaping
23x = torch.randn(2, 3, 4)
24reshaped = x.view(6, 4)        # Like NumPy reshape
25reshaped = x.reshape(6, 4)     # Alternative
26flat = x.flatten()              # Flatten to 1D
27print(f"Reshaped: {reshaped.shape}")
28print(f"Flattened: {flat.shape}")
29
30# Transpose
31matrix = torch.randn(3, 4)
32transposed = matrix.t()         # Transpose 2D
33print(f"Transposed shape: {transposed.shape}")
34
35# Permute dimensions (generalized transpose)
36x = torch.randn(2, 3, 4)
37permuted = x.permute(2, 0, 1)   # Move dims: [4, 2, 3]
38print(f"Permuted shape: {permuted.shape}")

1# Check GPU availability
2if torch.cuda.is_available():
3    device = torch.device("cuda")
4    print(f"GPU: {torch.cuda.get_device_name(0)}")
5else:
6    device = torch.device("cpu")
7    print("No GPU available, using CPU")
8
9# Move tensor to GPU
10x = torch.randn(1000, 1000)
11x_gpu = x.to(device)           # Move to GPU
12# Or: x_gpu = x.cuda()
13
14print(f"CPU tensor device: {x.device}")
15print(f"GPU tensor device: {x_gpu.device}")
16
17# Operations on GPU are automatic
18y_gpu = x_gpu * 2
19result_gpu = torch.mm(x_gpu, x_gpu.t())
20
21# Move back to CPU
22result_cpu = result_gpu.cpu()
23
24# Common pattern: device-agnostic code
25device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
26model = MyModel().to(device)
27data = data.to(device)

1x = torch.arange(0, 24).reshape(4, 6)
2
3# NumPy-like indexing
4print(x[0, :])       # First row
5print(x[:, 0])       # First column
6print(x[1:3, 2:5])   # Submatrix
7
8# Advanced indexing
9indices = torch.tensor([0, 2])
10selected = x[indices]  # Select specific rows
11
12# Boolean masking
13mask = x > 10
14filtered = x[mask]
15print(f"Values > 10: {filtered}")
16
17# In-place operations (end with _)
18x.add_(1)           # Add 1 in-place
19x.mul_(2)           # Multiply by 2 in-place
20x.zero_()           # Zero out in-place

1# Convert between PyTorch and NumPy
2tensor = torch.tensor([1, 2, 3])
3numpy_array = tensor.numpy()   # Tensor → NumPy (shares memory!)
4
5numpy_array = np.array([1, 2, 3])
6tensor = torch.from_numpy(numpy_array)  # NumPy → Tensor
7
8# Warning: They share memory!
9tensor = torch.ones(3)
10numpy_array = tensor.numpy()
11tensor.add_(1)  # Modifies tensor
12print(numpy_array)  # Also changed! [2. 2. 2.]
13
14# To avoid shared memory, use copy
15numpy_copy = tensor.numpy().copy()

1import torch
2
3# Create tensor with gradient tracking
4x = torch.tensor(2.0, requires_grad=True)
5print(f"x: {x}")
6print(f"x.grad: {x.grad}")  # None initially
7
8# Perform operations
9y = x ** 2 + 3 * x + 1
10print(f"y = x² + 3x + 1 = {y}")
11
12# Compute gradients
13y.backward()  # Compute dy/dx
14
15# Access gradient
16print(f"dy/dx = 2x + 3 = {x.grad}")  # At x=2: 2*2 + 3 = 7

1x = torch.tensor(2.0, requires_grad=True)
2w = torch.tensor(3.0, requires_grad=True)
3b = torch.tensor(1.0, requires_grad=True)
4
5# Forward pass
6y = w * x + b  # y = 3*2 + 1 = 7
7z = y ** 2     # z = 7² = 49
8
9print(f"Forward: z = {z}")
10
11# Backward pass
12z.backward()
13
14print(f"dz/dx = {x.grad}")  # 2*(wx+b)*w = 2*7*3 = 42
15print(f"dz/dw = {w.grad}")  # 2*(wx+b)*x = 2*7*2 = 28
16print(f"dz/db = {b.grad}")  # 2*(wx+b)*1 = 2*7*1 = 14
17
18# Computation graph (conceptual):
19# x, w, b (requires_grad=True)
20#    ↓
21# y = w*x + b (MulBackward, AddBackward)
22#    ↓
23# z = y**2 (PowBackward)
24#    ↓
25# z.backward() → propagates gradients backward

1# Gradients accumulate by default!
2x = torch.tensor(2.0, requires_grad=True)
3
4y = x ** 2
5y.backward()
6print(f"First backward: {x.grad}")  # 4
7
8y = x ** 2
9y.backward()
10print(f"Second backward: {x.grad}")  # 8 (accumulated!)
11
12# Must zero gradients manually
13x.grad.zero_()
14print(f"After zeroing: {x.grad}")  # 0

1x = torch.tensor(2.0, requires_grad=True)
2y = x ** 2
3
4# Detach: remove from computation graph
5y_detached = y.detach()
6print(f"y requires_grad: {y.requires_grad}")  # True
7print(f"y_detached requires_grad: {y_detached.requires_grad}")  # False
8
9# Use case: freezing parts of a model
10z = y_detached * 3  # This operation won't track gradients

1x = torch.randn(5, requires_grad=True)
2
3# During inference, disable gradient tracking
4with torch.no_grad():
5    y = x ** 2
6    z = y.mean()
7
8print(f"z requires_grad: {z.requires_grad}")  # False
9
10# Saves memory and speeds up inference!
11# Use this for validation/testing

1# Multi-input, multi-operation example
2x = torch.tensor(2.0, requires_grad=True)
3y = torch.tensor(3.0, requires_grad=True)
4
5# Complex function
6a = x * y           # 2 * 3 = 6
7b = x + y           # 2 + 3 = 5
8c = a ** 2          # 6² = 36
9d = b ** 3          # 5³ = 125
10output = c + d      # 36 + 125 = 161
11
12print(f"Output: {output}")
13
14# Compute all gradients with one call
15output.backward()
16
17print(f"∂output/∂x: {x.grad}")  # Chain rule applied!
18print(f"∂output/∂y: {y.grad}")
19
20# PyTorch automatically:
21# 1. Built computation graph during forward pass
22# 2. Applied chain rule during backward pass
23# 3. Computed all gradients efficiently

1class CustomSquare(torch.autograd.Function):
2    """
3    Custom operation with manual forward and backward
4    """
5    @staticmethod
6    def forward(ctx, input):
7        ctx.save_for_backward(input)  # Save for backward
8        return input ** 2
9
10    @staticmethod
11    def backward(ctx, grad_output):
12        input, = ctx.saved_tensors
13        grad_input = grad_output * 2 * input  # d(x²)/dx = 2x
14        return grad_input
15
16# Use custom function
17custom_square = CustomSquare.apply
18x = torch.tensor(3.0, requires_grad=True)
19y = custom_square(x)
20y.backward()
21print(f"Custom gradient: {x.grad}")  # 2*3 = 6

1import torch.nn as nn
2
3# Simple neural network
4model = nn.Linear(10, 1)  # 10 inputs → 1 output
5
6# Forward pass
7x = torch.randn(5, 10)  # Batch of 5 samples
8output = model(x)
9
10# Loss
11loss = output.mean()
12
13# Backward pass (autograd computes all gradients!)
14loss.backward()
15
16# Access gradients
17print(f"Weight gradient shape: {model.weight.grad.shape}")
18print(f"Bias gradient shape: {model.bias.grad.shape}")
19
20# Optimizer will use these gradients
21optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
22optimizer.step()  # Update weights: w -= lr * grad
23optimizer.zero_grad()  # Clear gradients for next iteration

1# 1. Always zero gradients before backward
2optimizer.zero_grad()
3loss.backward()
4optimizer.step()
5
6# 2. Use torch.no_grad() for inference
7with torch.no_grad():
8    predictions = model(test_data)
9
10# 3. Detach when you don't need gradients
11value = loss.item()  # Extract Python number (detaches)
12
13# 4. Handle gradient accumulation explicitly
14for i in range(accumulation_steps):
15    loss = compute_loss()
16    loss.backward()  # Accumulate gradients
17
18optimizer.step()
19optimizer.zero_grad()

1import torch
2import torch.nn as nn
3
4class SimpleNet(nn.Module):
5    def __init__(self, input_size, hidden_size, output_size):
6        super(SimpleNet, self).__init__()  # Call parent constructor
7
8        # Define layers
9        self.layer1 = nn.Linear(input_size, hidden_size)
10        self.relu = nn.ReLU()
11        self.layer2 = nn.Linear(hidden_size, output_size)
12
13    def forward(self, x):
14        """
15        Forward pass: define how data flows through the network
16        """
17        x = self.layer1(x)
18        x = self.relu(x)
19        x = self.layer2(x)
20        return x
21
22# Create model
23model = SimpleNet(input_size=10, hidden_size=20, output_size=2)
24
25# Forward pass
26input_data = torch.randn(5, 10)  # Batch of 5 samples
27output = model(input_data)
28print(f"Output shape: {output.shape}")  # [5, 2]
29
30# Access parameters
31print(f"Number of parameters: {sum(p.numel() for p in model.parameters())}")
32for name, param in model.named_parameters():
33    print(f"{name}: {param.shape}")

1# 1. Linear (Fully Connected)
2linear = nn.Linear(in_features=100, out_features=50)
3# Computes: y = xW^T + b
4
5# 2. Convolutional Layers
6conv2d = nn.Conv2d(in_channels=3, out_channels=64,
7                   kernel_size=3, padding=1)
8# For images: [batch, channels, height, width]
9
10# 3. Recurrent Layers
11lstm = nn.LSTM(input_size=50, hidden_size=100, num_layers=2)
12gru = nn.GRU(input_size=50, hidden_size=100)
13
14# 4. Activation Functions
15relu = nn.ReLU()
16sigmoid = nn.Sigmoid()
17tanh = nn.Tanh()
18softmax = nn.Softmax(dim=1)
19leaky_relu = nn.LeakyReLU(0.2)
20
21# 5. Normalization
22batch_norm = nn.BatchNorm1d(num_features=100)
23layer_norm = nn.LayerNorm(normalized_shape=100)
24
25# 6. Dropout (Regularization)
26dropout = nn.Dropout(p=0.5)  # 50% dropout rate
27
28# 7. Pooling
29max_pool = nn.MaxPool2d(kernel_size=2, stride=2)
30avg_pool = nn.AvgPool2d(kernel_size=2, stride=2)

1# Simple way to stack layers
2model = nn.Sequential(
3    nn.Linear(784, 256),
4    nn.ReLU(),
5    nn.Dropout(0.2),
6    nn.Linear(256, 128),
7    nn.ReLU(),
8    nn.Dropout(0.2),
9    nn.Linear(128, 10),
10    nn.Softmax(dim=1)
11)
12
13# Equivalent to:
14class SequentialNet(nn.Module):
15    def __init__(self):
16        super().__init__()
17        self.layer1 = nn.Linear(784, 256)
18        self.relu1 = nn.ReLU()
19        self.dropout1 = nn.Dropout(0.2)
20        # ... etc
21
22    def forward(self, x):
23        x = self.layer1(x)
24        x = self.relu1(x)
25        x = self.dropout1(x)
26        # ... etc
27        return x

1class ConvNet(nn.Module):
2    def __init__(self, num_classes=10):
3        super(ConvNet, self).__init__()
4
5        # Convolutional blocks
6        self.conv_block1 = nn.Sequential(
7            nn.Conv2d(3, 32, kernel_size=3, padding=1),
8            nn.BatchNorm2d(32),
9            nn.ReLU(inplace=True),
10            nn.MaxPool2d(kernel_size=2, stride=2)
11        )
12
13        self.conv_block2 = nn.Sequential(
14            nn.Conv2d(32, 64, kernel_size=3, padding=1),
15            nn.BatchNorm2d(64),
16            nn.ReLU(inplace=True),
17            nn.MaxPool2d(kernel_size=2, stride=2)
18        )
19
20        # Fully connected layers
21        self.classifier = nn.Sequential(
22            nn.Dropout(0.5),
23            nn.Linear(64 * 8 * 8, 512),
24            nn.ReLU(inplace=True),
25            nn.Dropout(0.5),
26            nn.Linear(512, num_classes)
27        )
28
29    def forward(self, x):
30        # Input: [batch, 3, 32, 32]
31        x = self.conv_block1(x)  # [batch, 32, 16, 16]
32        x = self.conv_block2(x)  # [batch, 64, 8, 8]
33
34        # Flatten
35        x = x.view(x.size(0), -1)  # [batch, 64*8*8]
36
37        # Classify
38        x = self.classifier(x)  # [batch, num_classes]
39        return x
40
41model = ConvNet(num_classes=10)
42print(model)

1class DynamicNet(nn.Module):
2    def __init__(self):
3        super().__init__()
4        self.layer1 = nn.Linear(10, 20)
5        self.layer2 = nn.Linear(20, 10)
6        self.relu = nn.ReLU()
7
8    def forward(self, x, num_iterations=3):
9        """
10        Dynamic computation: different graph each forward pass
11        """
12        x = self.layer1(x)
13
14        # Dynamic loop (number of iterations can change!)
15        for _ in range(num_iterations):
16            x = self.relu(x)
17            x = self.layer2(x)
18
19        return x
20
21model = DynamicNet()
22output1 = model(torch.randn(5, 10), num_iterations=2)
23output2 = model(torch.randn(5, 10), num_iterations=5)
24# Different computation graphs!

1# Move model to GPU
2device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
3model = model.to(device)
4
5# Set training/evaluation mode
6model.train()   # Enable dropout, batch norm training mode
7model.eval()    # Disable dropout, batch norm eval mode
8
9# Save and load models
10torch.save(model.state_dict(), 'model.pth')
11model.load_state_dict(torch.load('model.pth'))
12
13# Save entire model
14torch.save(model, 'complete_model.pth')
15model = torch.load('complete_model.pth')
16
17# Freeze layers (transfer learning)
18for param in model.conv_block1.parameters():
19    param.requires_grad = False
20
21# Count parameters
22total_params = sum(p.numel() for p in model.parameters())
23trainable_params = sum(p.numel() for p in model.parameters()
24                       if p.requires_grad)
25
26print(f"Total: {total_params:,}, Trainable: {trainable_params:,}")

1def init_weights(m):
2    if isinstance(m, nn.Linear):
3        nn.init.xavier_uniform_(m.weight)
4        nn.init.zeros_(m.bias)
5    elif isinstance(m, nn.Conv2d):
6        nn.init.kaiming_normal_(m.weight, mode='fan_out',
7                                nonlinearity='relu')
8
9model.apply(init_weights)  # Apply to all layers

1import torch
2from torch.utils.data import Dataset, DataLoader
3
4# 1. Custom Dataset
5class CustomDataset(Dataset):
6    """
7    Dataset must implement:
8    - __len__: return number of samples
9    - __getitem__: return one sample
10    """
11    def __init__(self, data, labels):
12        self.data = data
13        self.labels = labels
14
15    def __len__(self):
16        return len(self.data)
17
18    def __getitem__(self, idx):
19        return self.data[idx], self.labels[idx]
20
21# Create dataset
22X = torch.randn(1000, 10)  # 1000 samples, 10 features
23y = torch.randint(0, 2, (1000,))  # Binary labels
24
25dataset = CustomDataset(X, y)
26
27# 2. DataLoader: handles batching, shuffling, parallel loading
28dataloader = DataLoader(
29    dataset,
30    batch_size=32,        # Process 32 samples at once
31    shuffle=True,         # Shuffle each epoch
32    num_workers=0,        # Parallel data loading (set > 0 on Linux/Mac)
33    pin_memory=True       # Speed up GPU transfer
34)
35
36# Iterate through batches
37for batch_idx, (data, labels) in enumerate(dataloader):
38    print(f"Batch {batch_idx}: data shape {data.shape}, labels shape {labels.shape}")
39    if batch_idx >= 2:
40        break

1import torch
2import torch.nn as nn
3import torch.optim as optim
4from torch.utils.data import DataLoader, TensorDataset
5
6# 1. Prepare Data
7X_train = torch.randn(1000, 20)
8y_train = torch.randint(0, 10, (1000,))
9
10train_dataset = TensorDataset(X_train, y_train)
11train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
12
13# 2. Define Model
14model = nn.Sequential(
15    nn.Linear(20, 64),
16    nn.ReLU(),
17    nn.Dropout(0.2),
18    nn.Linear(64, 32),
19    nn.ReLU(),
20    nn.Linear(32, 10)
21)
22
23# 3. Loss and Optimizer
24criterion = nn.CrossEntropyLoss()
25optimizer = optim.Adam(model.parameters(), lr=0.001)
26
27# 4. Training Loop
28device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
29model.to(device)
30
31num_epochs = 10
32
33for epoch in range(num_epochs):
34    model.train()  # Set to training mode
35
36    running_loss = 0.0
37    correct = 0
38    total = 0
39
40    for batch_idx, (data, labels) in enumerate(train_loader):
41        # Move to device
42        data, labels = data.to(device), labels.to(device)
43
44        # Zero gradients
45        optimizer.zero_grad()
46
47        # Forward pass
48        outputs = model(data)
49        loss = criterion(outputs, labels)
50
51        # Backward pass
52        loss.backward()
53
54        # Update weights
55        optimizer.step()
56
57        # Statistics
58        running_loss += loss.item()
59        _, predicted = torch.max(outputs.data, 1)
60        total += labels.size(0)
61        correct += (predicted == labels).sum().item()
62
63    # Epoch statistics
64    epoch_loss = running_loss / len(train_loader)
65    epoch_acc = 100 * correct / total
66
67    print(f"Epoch [{epoch+1}/{num_epochs}] "
68          f"Loss: {epoch_loss:.4f} "
69          f"Accuracy: {epoch_acc:.2f}%")

1# Prepare validation data
2X_val = torch.randn(200, 20)
3y_val = torch.randint(0, 10, (200,))
4val_dataset = TensorDataset(X_val, y_val)
5val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
6
7def validate(model, val_loader, criterion, device):
8    model.eval()  # Set to evaluation mode
9
10    val_loss = 0.0
11    correct = 0
12    total = 0
13
14    with torch.no_grad():  # Disable gradient computation
15        for data, labels in val_loader:
16            data, labels = data.to(device), labels.to(device)
17
18            outputs = model(data)
19            loss = criterion(outputs, labels)
20
21            val_loss += loss.item()
22            _, predicted = torch.max(outputs.data, 1)
23            total += labels.size(0)
24            correct += (predicted == labels).sum().item()
25
26    avg_loss = val_loss / len(val_loader)
27    accuracy = 100 * correct / total
28
29    return avg_loss, accuracy
30
31# Call during training
32val_loss, val_acc = validate(model, val_loader, criterion, device)
33print(f"Validation Loss: {val_loss:.4f}, Accuracy: {val_acc:.2f}%")

1# Adjust learning rate during training
2scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
3# Reduces LR by 10x every 5 epochs
4
5for epoch in range(num_epochs):
6    # Training code...
7
8    scheduler.step()  # Update learning rate
9    current_lr = optimizer.param_groups[0]['lr']
10    print(f"Epoch {epoch+1}, LR: {current_lr}")

1class EarlyStopping:
2    def __init__(self, patience=5, min_delta=0):
3        self.patience = patience
4        self.min_delta = min_delta
5        self.counter = 0
6        self.best_loss = None
7        self.early_stop = False
8
9    def __call__(self, val_loss):
10        if self.best_loss is None:
11            self.best_loss = val_loss
12        elif val_loss > self.best_loss - self.min_delta:
13            self.counter += 1
14            if self.counter >= self.patience:
15                self.early_stop = True
16        else:
17            self.best_loss = val_loss
18            self.counter = 0
19
20early_stopping = EarlyStopping(patience=3)
21
22for epoch in range(num_epochs):
23    # Training...
24    val_loss, val_acc = validate(model, val_loader, criterion, device)
25
26    early_stopping(val_loss)
27    if early_stopping.early_stop:
28        print("Early stopping triggered")
29        break

1def train_model(model, train_loader, val_loader, criterion,
2                optimizer, num_epochs, device):
3    """
4    Complete training function with validation and checkpointing
5    """
6    best_val_acc = 0.0
7
8    for epoch in range(num_epochs):
9        # Training phase
10        model.train()
11        train_loss = 0.0
12
13        for data, labels in train_loader:
14            data, labels = data.to(device), labels.to(device)
15
16            optimizer.zero_grad()
17            outputs = model(data)
18            loss = criterion(outputs, labels)
19            loss.backward()
20            optimizer.step()
21
22            train_loss += loss.item()
23
24        # Validation phase
25        val_loss, val_acc = validate(model, val_loader, criterion, device)
26
27        print(f"Epoch [{epoch+1}/{num_epochs}]")
28        print(f"  Train Loss: {train_loss/len(train_loader):.4f}")
29        print(f"  Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%")
30
31        # Save best model
32        if val_acc > best_val_acc:
33            best_val_acc = val_acc
34            torch.save(model.state_dict(), 'best_model.pth')
35            print(f"  Saved new best model!")
36
37    return model
38
39# Usage
40trained_model = train_model(model, train_loader, val_loader,
41                            criterion, optimizer, num_epochs=20,
42                            device=device)