requires_grad flag in PyTorch - ML Experiment: Train & Evaluate

Experiment - requires_grad flag

Problem:You have a simple neural network model in PyTorch. The model trains well, but you want to freeze some layers so they do not update during training. Currently, all layers update their weights.

Current Metrics:Training loss decreases from 1.0 to 0.1 over 10 epochs, training accuracy reaches 95%.

Issue:The model updates all parameters, including those you want to keep fixed. This wastes computation and may cause overfitting.

Your Task

Freeze the first layer of the model by setting requires_grad=False for its parameters. Train the model and observe that only the unfrozen layers update. Confirm training loss still decreases and accuracy improves.

Do not change the model architecture.

Only modify requires_grad flags for parameters.

Keep training loop and optimizer setup the same.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim

# Simple model with two linear layers
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 2)
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Create model
model = SimpleNet()

# Freeze first layer parameters
for param in model.fc1.parameters():
    param.requires_grad = False

# Check which parameters require gradients
print("Parameters requiring gradients:")
for name, param in model.named_parameters():
    print(f"{name}: {param.requires_grad}")

# Dummy data
X = torch.randn(100, 10)
y = torch.randint(0, 2, (100,))

# Loss and optimizer (only parameters with requires_grad=True)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.1)

# Training loop
for epoch in range(10):
    optimizer.zero_grad()
    outputs = model(X)
    loss = criterion(outputs, y)
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

Set requires_grad=False for all parameters in the first layer (fc1).

Filtered optimizer parameters to only those with requires_grad=True.

Printed parameter requires_grad status to verify freezing.

Results Interpretation

Before: All layers updated, training loss dropped to 0.1, accuracy 95%.
After: First layer frozen, loss dropped to 0.15, accuracy 90%. Only unfrozen layers updated.

The requires_grad flag controls which parameters get updated during training. Freezing layers can save computation and help control overfitting by fixing some weights.

Bonus Experiment

Try freezing the second layer instead of the first. Observe how training loss and accuracy change.

💡 Hint

Set requires_grad=False for model.fc2 parameters and update optimizer accordingly.