PyTorch Interview Questions: Training Loop

PyTorch/training-loop-pytorch

In a standard PyTorch training loop, if you do not call optimizer.zero_grad() before loss.backward(), gradients from the previous batch will accumulate (sum) into the same .grad tensors instead of being replaced.#

PyTorch/training-loop-pytorch/loss-functions-pytorch

In a PyTorch training loop, you are training a classifier whose final layer is a nn.Linear with out_features=num_classes (raw logits, no softmax applied). Which loss function is the standard, correct choice for multi-class classification with integer class labels?#

PyTorch/training-loop-pytorch/loss-functions-pytorch

In a standard PyTorch training loop, after computing the loss tensor with loss = criterion(outputs, labels), the next two steps are to call loss._____() to populate .grad on all leaf tensors that require gradients, then call optimizer._____() to update the model parameters using those gradients. (Before the next iteration you would also call optimizer.zero_grad().)#

PyTorch/training-loop-pytorch

Given the following PyTorch training snippet, what is printed to stdout? Assume all imports are present and torch seeds are irrelevant because weights are set explicitly.#

import torch
import torch.nn as nn

model = nn.Linear(2, 1, bias=False)
model.weight.data = torch.tensor([[2.0, 2.0]])

criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

x = torch.tensor([[1.0, 1.0]])
y = torch.tensor([[0.0]])

for epoch in range(3):
    optimizer.zero_grad()
    out = model(x)
    loss = criterion(out, y)
    loss.backward()
    optimizer.step()
    print(round(model.weight[0, 0].item(), 4))

1.2
0.72
0.432

PyTorch/training-loop-pytorch

The following PyTorch training loop is intended to update model weights using per-batch gradients. It contains a single bug. Identify the buggy line number (1-based).#

import torch
import torch.nn as nn

model = nn.Linear(10, 1)
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for inputs, labels in dataloader:
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    optimizer.step()
    loss.backward()

PyTorch/training-loop-pytorch/loss-functions-pytorch

What value is printed by the following PyTorch snippet? Round your answer to 4 decimal places.#

import torch
import torch.nn as nn

logits = torch.tensor([[2.0, 1.0, 0.1]])
target = torch.tensor([0])
criterion = nn.CrossEntropyLoss()
loss = criterion(logits, target)
print(round(loss.item(), 4))

0.4168

PyTorch/training-loop-pytorch/loss-functions-pytorch

In a PyTorch training loop for a 3-class classifier, your model's final linear layer outputs raw logits (no softmax applied). Which built-in loss function should you use so that gradients are computed correctly without manually applying softmax?#

PyTorch/training-loop-pytorch

Trace the following PyTorch training loop that implements SGD with momentum manually (without using torch.optim). The code runs on CPU with two training examples processed sequentially. What is the exact printed output?#

import torch

torch.manual_seed(42)

w = torch.tensor(1.0, requires_grad=True)
b = torch.tensor(0.0, requires_grad=True)

lr = 0.1
momentum = 0.9

v_w = torch.tensor(0.0)
v_b = torch.tensor(0.0)

data = [
    (torch.tensor(2.0), torch.tensor(5.0)),
    (torch.tensor(3.0), torch.tensor(7.0)),
]

for x, y_true in data:
    y_pred = w * x + b
    loss = (y_pred - y_true) ** 2
    loss.backward()

    with torch.no_grad():
        v_w = momentum * v_w + w.grad
        v_b = momentum * v_b + b.grad
        w -= lr * v_w
        b -= lr * v_b
        w.grad.zero_()
        b.grad.zero_()

    print(f"w={w.item():.4f}, b={b.item():.4f}")

w=2.2000, b=0.6000
w=3.1600, b=1.1000

PyTorch/training-loop-pytorch

The following PyTorch training loop is intended to train a classifier with BatchNorm and Dropout for 50 epochs, validating each epoch. The model trains (loss decreases) but performs poorly and never generalizes. Identify the buggy line.#

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.bn1 = nn.BatchNorm1d(256)
        self.drop = nn.Dropout(0.5)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.bn1(x)
        x = self.drop(x)
        return self.fc2(x)

model = Net()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

for epoch in range(50):
    model.eval()
    for x, y in train_loader:
        optimizer.zero_grad()
        out = model(x)
        loss = criterion(out, y)
        loss.backward()
        optimizer.step()

    model.eval()
    with torch.no_grad():
        for x, y in val_loader:
            val_loss = criterion(model(x), y)

PyTorch/training-loop-pytorch/loss-functions-pytorch

The following training loop compiles and runs without error, but the model trains poorly — the loss decreases very slowly and accuracy plateaus well below what the architecture should achieve. Find the bug.#

import torch
import torch.nn as nn
import torch.nn.functional as F

class Classifier(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.fc = nn.Linear(256, num_classes)

    def forward(self, x):
        x = self.fc(x)
        return F.softmax(x, dim=1)

model = Classifier(10)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for inputs, targets in train_loader:
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()