Deep Learning Frameworks

Question 1 (12 marks)

(a) y = w @ x:

• Row 0: $1(2) + (-1)(3) = -1$
• Row 1: $0(2) + 2(3) = 6$

So $y = [-1, 6]$. (1.5) $z = (-1)^2 + 6^2 = 1 + 36 = 37$. (1.5)

(b) $z = \sum_i y_i^2$ where $y = Wx$. So $\dfrac{\partial z}{\partial x} = 2 W^\top y$. (chain rule: $\partial z/\partial y = 2y$, and $\partial y/\partial x = W$, so $\partial z/\partial x = W^\top (2y)$.) (3)

Numeric: $2y = [-2, 12]$. $W^\top = \begin{bmatrix}1 & 0\\-1 & 2\end{bmatrix}$. $W^\top(2y) = [1(-2)+0(12),\ -1(-2)+2(12)] = [-2,\ 26]$. So x.grad = [-2., 26.]. (2)

(c) $\dfrac{\partial z}{\partial W_{ij}} = 2y_i x_j$, i.e. outer product $2y\, x^\top$. (2) $2y = [-2,12]^\top$, $x=[2,3]$: $w.grad = \begin{bmatrix}-2\cdot2 & -2\cdot3\\ 12\cdot2 & 12\cdot3\end{bmatrix} = \begin{bmatrix}-4 & -6\\ 24 & 36\end{bmatrix}$ (1) Shape matches w because autograd stores a gradient of equal shape for every parameter tensor (one derivative per element). (1)

Question 2 (12 marks)

(a) Bug 1 (runtime error): missing super().__init__() — submodules aren't registered, .parameters() empty / attribute errors. Bug 2 (learning capacity): no non-linearity between layers, so the stack collapses to a single linear map. (3)

Corrected: (3)

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 32)
        self.fc2 = nn.Linear(32, 1)
    def forward(self, x):
        x = torch.relu(self.fc1(x))   # non-linearity added
        x = self.fc2(x)
        return x

(b) Output shape for input (64,10): (64, 1). (1) Params: fc1 = $10\times32 + 32 = 352$; fc2 = $32\times1 + 1 = 33$. Total $= 385$. (3)

(c) model.eval() sets modules into evaluation mode; good habit for reproducibility. It only changes behaviour if the model contains mode-dependent layers (Dropout, BatchNorm, etc.) — here it makes no functional difference, but is correct practice. (2)

Question 3 (14 marks)

(a) (6) (1 each: DataLoader construction, loop, to(device), zero_grad, forward+loss, backward+step)

loader = torch.utils.data.DataLoader(train_ds, batch_size=128, shuffle=True)
model.train()
for xb, yb in loader:
    xb, yb = xb.to(device), yb.to(device)
    opt.zero_grad()
    out = model(xb)
    loss = criterion(out, yb)
    loss.backward()
    opt.step()

(b) $\lceil 50000/128\rceil = \lceil 390.625\rceil = 391$ steps. (2)

(c) Without zero_grad, .backward() accumulates gradients into .grad across batches (they sum rather than reset). Effective gradient magnitude grows batch-over-batch, producing oversized/erratic parameter updates → divergence or NaN loss. (3)

(d) Increase the learning rate (roughly proportionally, e.g. ~4× for 4× batch, i.e. "linear scaling rule"). Larger batches give lower-variance gradient estimates, so a larger step is stable/needed to keep convergence speed comparable. (Alt acceptable: add warmup.) (3)

Question 4 (12 marks)

(a) (6)

scaler = torch.cuda.amp.GradScaler()
opt.zero_grad()
with torch.cuda.amp.autocast():
    out = model(xb)
    loss = criterion(out, yb)
scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()

Marks: autocast context (2), scale().backward() (2), step+update order (2).

(b) fp16 has a narrow exponent range; small gradients underflow to zero, killing updates. GradScaler multiplies the loss by a large factor before backward so gradients stay representable, then unscales them before the optimizer step, restoring true magnitude and skipping steps if inf/nan detected. (3)

(c) (3)

model = model.to(device)
xb, yb = xb.to(device), yb.to(device)

Root cause: an operation mixed a tensor on GPU with one on CPU (e.g. inputs not moved, or a buffer/parameter left on CPU); all operands of an op must share one device.

Question 5 (10 marks)

(a) state_dict stores only tensor weights → portable, decoupled from code. Saving the whole object pickles the class definition/path; failure mode: loading breaks (or silently misbehaves) if the class file moved, was renamed, or the library version changed. (3)

(b) (4)

# save
torch.save({'epoch': epoch,
            'model_state': model.state_dict(),
            'opt_state': opt.state_dict()}, 'ckpt.pt')
# load
ckpt = torch.load('ckpt.pt')
model.load_state_dict(ckpt['model_state'])
opt.load_state_dict(ckpt['opt_state'])
start_epoch = ckpt['epoch'] + 1

(c) Global batch size $= 4 \times 64 = 256$. (2) Gradients are synchronized by an all-reduce operation (averaging gradients across processes during backward). (1)

[
  {"claim":"Q1 z value equals 37","code":"x=Matrix([2,3]);W=Matrix([[1,-1],[0,2]]);y=W*x;z=(y.T*y)[0];result=(z==37)"},
  {"claim":"Q1 x.grad equals [-2,26]","code":"x=Matrix([2,3]);W=Matrix([[1,-1],[0,2]]);y=W*x;g=2*W.T*y;result=(list(g)==[-2,26])"},
  {"claim":"Q1 w.grad equals [[-4,-6],[24,36]]","code":"x=Matrix([2,3]);y=Matrix([-1,6]);G=2*y*x.T;result=(G==Matrix([[-4,-6],[24,36]]))"},
  {"claim":"Q2 total params equal 385","code":"result=((10*32+32)+(32*1+1)==385)"},
  {"claim":"Q3 optimizer steps equal 391","code":"import math;result=(math.ceil(50000/128)==391)"},
  {"claim":"Q5 global batch size equals 256","code":"result=(4*64==256)"}
]