Minibatches

Minibatches

GD: \mathcal{O}(n) per step, exact. SGD: \mathcal{O}(1) per step, noisy. Everything we ever trained sat in between — a minibatch of b examples:

\mathbf{w} \leftarrow \mathbf{w} - \frac{\eta}{b} \sum_{i \in \mathcal{B}} \nabla f_i(\mathbf{w}).

• Statistics (last section): variance \propto 1/b.
• This section: the mechanics — why b at once costs far less than b one at a time.

Arithmetic outruns memory

• Server CPU: 10^{12}–10^{13} FLOP/s vs a few 100 GB/s of bandwidth.
• GPU: 10^{14+} FLOP/s vs a few TB/s.
• Ratio ≈ two orders of magnitude: each byte loaded must feed tens to hundreds of operations.

Caches bridge the gap only if the algorithm reuses resident data — blocked matmul, not elementwise loops.

Plus dispatch overhead: every op launched from Python costs microseconds; the arithmetic inside a tiny op costs nanoseconds.

Setup

%matplotlib inline
from d2l import torch as d2l
import numpy as np
import time
import torch
from torch import nn

A = torch.zeros(256, 256)
B = torch.randn(256, 256)
C = torch.randn(256, 256)

class Timer:
    """Record multiple running times."""
    def __init__(self):
        self.times = []
        self.start()

    def start(self):
        """Start the timer."""
        self.tik = time.time()

    def stop(self):
        """Stop the timer and record the time in a list."""
        self.times.append(time.time() - self.tik)
        return self.times[-1]

    def avg(self):
        """Return the average time."""
        return sum(self.times) / len(self.times)

    def sum(self):
        """Return the sum of time."""
        return sum(self.times)

    def cumsum(self):
        """Return the accumulated time."""
        return np.array(self.times).cumsum().tolist()

timer = Timer()

Three loops, three speeds

Compute \mathbf{A} = \mathbf{B}\mathbf{C} on 256 \times 256 matrices, exposing more work per call each time:

# Compute A = BC one element at a time
timer.start()
for i in range(256):
    for j in range(256):
        A[i, j] = torch.dot(B[i, :], C[:, j])
timer.stop()

0.8051884174346924

# Compute A = BC one column at a time
timer.start()
for j in range(256):
    A[:, j] = torch.mv(B, C[:, j])
timer.stop()

0.013139963150024414

# Compute A = BC in one go
timer.start()
A = torch.mm(B, C)
timer.stop()

gigaflops = [0.03 / i for i in timer.times]
print(f'performance in Gigaflops: element {gigaflops[0]:.3f}, '
      f'column {gigaflops[1]:.3f}, full {gigaflops[2]:.3f}')

performance in Gigaflops: element 0.037, column 2.283, full 20.689

Same arithmetic, orders of magnitude apart. The loop is overhead; the cache and vector units do the work.

Batching in blocks

64 columns at a time — a “minibatch” of the matmul:

timer.start()
for j in range(0, 256, 64):
    A[:, j:j+64] = torch.mm(B, C[:, j:j+64])
timer.stop()
print(f'performance in Gigaflops: block {0.03 / timer.times[3]:.3f}')

performance in Gigaflops: block 17.145

Already as fast as the full multiplication: modest batches amortize essentially all the overhead.

• Two reasons to batch, kept apart: hardware (this section) and variance (last section). Both saturate.
• How large is too large → critical batch size, later in this chapter.
• Batch exceeds memory? Gradient accumulation (ch. on performance).

Airfoil dataset

Real regression data for the whole chapter — 1500 examples, 5 features, whitened; every run takes seconds:

d2l.DATA_HUB['airfoil'] = (d2l.DATA_URL + 'airfoil_self_noise.dat',
                           '76e5be1548fd8222e5074cf0faae75edff8cf93f')


def get_data_ch11(batch_size=10, n=1500):
    data = np.genfromtxt(d2l.download('airfoil'),
                         dtype=np.float32, delimiter='\t')
    data = torch.from_numpy((data - data.mean(axis=0)) / data.std(axis=0))
    data_iter = d2l.load_array((data[:n, :-1], data[:n, -1]),
                               batch_size, is_train=True)
    return data_iter, data.shape[1]-1

The optimizer interface

Every optimizer in this chapter: (params, states, hyperparams) — states carries the algorithm’s memory (empty for SGD):

def sgd(params, states, hyperparams):
    for p in params:
        with torch.no_grad():
            p.sub_(hyperparams['lr'] * p.grad)
        p.grad.zero_()

Generic training harness

Linear regression + any update rule; loss recorded against wall-clock time:

def train_ch11(trainer_fn, states, hyperparams, data_iter,
               feature_dim, num_epochs=2):
    # Initialization
    w = torch.normal(mean=0.0, std=0.01, size=(feature_dim, 1),
                     requires_grad=True)
    b = torch.zeros((1), requires_grad=True)
    net, loss = lambda X: d2l.linreg(X, w, b), d2l.squared_loss
    # Train
    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                            xlim=[0, num_epochs], ylim=[0.22, 0.35])
    n, timer = 0, d2l.Timer()
    for _ in range(num_epochs):
        for X, y in data_iter:
            l = loss(net(X), y).mean()
            l.backward()
            trainer_fn([w, b], states, hyperparams)
            n += X.shape[0]
            if n % 200 == 0:
                timer.stop()
                animator.add(n/X.shape[0]/len(data_iter),
                             (d2l.evaluate_loss(net, data_iter, loss),))
                timer.start()
    print(f'loss: {animator.Y[0][-1]:.3f}, {timer.sum()/num_epochs:.3f} sec/epoch')
    return timer.cumsum(), animator.Y[0]

The race: full batch vs single example

b = 1500: one well-aimed update per epoch — stalls after ~6 steps:

def train_sgd(lr, batch_size, num_epochs=2):
    data_iter, feature_dim = get_data_ch11(batch_size)
    return train_ch11(
        sgd, None, {'lr': lr}, data_iter, feature_dim, num_epochs)

gd_res = train_sgd(1, 1500, 10)

loss: 0.244, 0.029 sec/epoch

b = 1: 1500 updates per epoch, but 1500 tiny dispatches — more clock time per epoch than GD:

sgd_res = train_sgd(0.005, 1)

loss: 0.245, 0.399 sec/epoch

The middle wins

mini1_res = train_sgd(.4, 100)

loss: 0.244, 0.018 sec/epoch

mini2_res = train_sgd(.05, 10)

loss: 0.246, 0.056 sec/epoch

d2l.set_figsize([6, 3])
d2l.plot(*list(map(list, zip(gd_res, sgd_res, mini1_res, mini2_res))),
         'time (sec)', 'loss', xlim=[1e-2, 10], xscale='log',
         legend=['gd', 'sgd', 'batch size=100', 'batch size=10'])

Read the x-axis as elapsed seconds: b=100 beats both extremes.

Concise: framework optimizer

Same experiment through the built-in optimizer — the harness the rest of the chapter reuses:

def train_concise_ch11(trainer_fn, hyperparams, data_iter, num_epochs=4):
    # Initialization
    net = nn.Sequential(nn.Linear(5, 1))
    def init_weights(module):
        if type(module) == nn.Linear:
            torch.nn.init.normal_(module.weight, std=0.01)
    net.apply(init_weights)

    optimizer = trainer_fn(net.parameters(), **hyperparams)
    loss = nn.MSELoss(reduction='none')
    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                            xlim=[0, num_epochs], ylim=[0.22, 0.35])
    n, timer = 0, d2l.Timer()
    for _ in range(num_epochs):
        for X, y in data_iter:
            optimizer.zero_grad()
            out = net(X)
            y = y.reshape(out.shape)
            l = loss(out, y)
            l.mean().backward()
            optimizer.step()
            n += X.shape[0]
            if n % 200 == 0:
                timer.stop()
                # `MSELoss` computes squared error without the 1/2 factor
                animator.add(n/X.shape[0]/len(data_iter),
                             (d2l.evaluate_loss(net, data_iter, loss) / 2,))
                timer.start()
    print(f'loss: {animator.Y[0][-1]:.3f}, {timer.sum()/num_epochs:.3f} sec/epoch')

data_iter, _ = get_data_ch11(10)
trainer = torch.optim.SGD
train_concise_ch11(trainer, {'lr': 0.01}, data_iter)

loss: 0.242, 0.059 sec/epoch

Recap

• Minibatch SGD interpolates between GD and SGD and beats both on the wall clock.
• The win is mechanical: dispatch amortization, cache reuse, vector units — plus the 1/b variance cut from last section.
• Pick b to fill the accelerator within memory; the statistical ceiling (critical batch size) comes later in the chapter.