Reading data, and a program that learns digits

url= and sha256=

%xs = buffer 3920000 dtype=f32 slot=0 url=data/mnist-train-images.f32 sha256=1ac7f6c8...

The buffer's contents come from the file instead of a literal list. The format is the simplest there is: the elements back to back, little-endian (each element's lowest byte first), no header. That's 4 bytes per f32 or i32 and 1 byte per bool, which must be 0 or 1. A file of the wrong size is an error, and so is a hash that doesn't match.

SHA-256 turns any file into a 64-hex-digit fingerprint, and changing one byte changes the fingerprint. The hash is what makes a URL safe to depend on. Files on the internet change, mirrors go stale, and downloads get cut off. With sha256=, a program means exactly one set of bytes: if what you fetched isn't it, the program is invalid rather than quietly wrong. It also means a download can be cached by its hash and never fetched twice.

url= is either https://... or a relative path, read from the folder the .uir file is in. Relative paths may not contain .., so a program can't reach outside its own folder. The test suite only uses relative paths, and the grader copies each data file next to the program it runs, so your compiler never needs the network to pass. For https, fetching with curl into a cache folder is plenty.

Where the bytes go

In run, your compiler reads the file, checks it, and decodes it into the buffer's values. In render, nothing changes: the C harness fills every buffer before calling your kernel, file buffers included, with one fread straight into a float array. Little-endian is the byte order of x86 and ARM, so no conversion is needed. That's why the format is what it is.

The data

MNIST is 70,000 handwritten digits, each a 28×28 grayscale image with a label from 0 to 9. The course uses a fixed slice: the first 5000 training images and the first 1000 test images. tools/mnist.py converted them from the official files (checking their published hashes) into the raw format above, with pixels as f32 from 0 to 1. The four files live in tests/week6/render/data/: 15.7 MB of training pixels, 3.1 MB of test pixels, and two small label files.

The model

The simplest classifier that works: softmax regression. One weight matrix W, 784×10, starting at zero. For an image x (784 pixels), the ten logits z = x W are scores, one per digit. Softmax turns them into probabilities, and the loss for one image is -log p[label]: small when the model puts high probability on the right digit. Averaged over a batch of 100, that's one number, and its gradient with respect to W says how to change every weight.

A step is W ← W - 1.0 · gradient(loss, W). The 1.0 is the learning rate: how far each step moves. Fifty steps of 100 images go once through the 5000 training digits.

One step, line by line

This is step 0 of tests/week6/render/mnist_train.uir, which python tools/mnist.py program writes:

# the batch: 100 images, and their labels as one-hot rows
%s0_xv = shrink %xs 0 shape=[78400]
%s0_xl = load %s0_xv
%s0_x = reshape %s0_xl shape=[100 784]
%s0_yv = shrink %ys 0 shape=[100]
%s0_yl = load %s0_yv
%s0_ohl = reshape %s0_yl shape=[100 1]
%s0_ohe = expand %s0_ohl shape=[100 10]
%s0_ohc = reshape %classes shape=[1 10]
%s0_ohk = expand %s0_ohc shape=[100 10]
%s0_ohq = cmpeq %s0_ohe %s0_ohk          # label == class, for every pair
%s0_oh = cast %s0_ohq dtype=f32
# the logits, staged: computed once, and so is their gradient
%s0_wl = load %w
%s0_W = reshape %s0_wl shape=[784 10]
%s0_z0 = matmul %s0_x %s0_W
%s0_z = stage %s0_z0
# log-softmax, made safe by subtracting each row's max (detached)
%s0_zt = permute %s0_z order=[1 0]
%s0_m = reduce %s0_zt op=max
%s0_md = detach %s0_m
%s0_mr = reshape %s0_md shape=[100 1]
%s0_me = expand %s0_mr shape=[100 10]
%s0_zs = sub %s0_z %s0_me
%s0_e = exp %s0_zs
%s0_et = permute %s0_e order=[1 0]
%s0_se = reduce %s0_et op=add
%s0_ls = log %s0_se
%s0_lr = reshape %s0_ls shape=[100 1]
%s0_le = expand %s0_lr shape=[100 10]
%s0_lp = sub %s0_zs %s0_le
# the loss: minus the mean log-probability of the right digit
%s0_pick = mul %s0_lp %s0_oh
%s0_pf = reshape %s0_pick shape=[1000]
%s0_tot = reduce %s0_pf op=add
%s0_loss = mul %s0_tot -0.01
# the update, stored back into the weights
%s0_g = gradient %s0_loss %s0_W
%s0_lrx = expand 1.0 shape=[784 10]
%s0_d = mul %s0_g %s0_lrx
%s0_W1 = sub %s0_W %s0_d
%s0_wf = reshape %s0_W1 shape=[7840]
%s0_st = store %w %s0_wf
%s0_w = after %w %s0_st

A few lines deserve a second look.

• The max subtraction. exp overflows f32 above about 88, so the program subtracts each row's largest logit first. The answer doesn't change, because softmax is the same if you shift every score by a constant. The shift's gradient is zero in exact math, so detach skips computing it.
• The one-hot rows. Each label becomes ten numbers: 1.0 at the true digit and 0.0 elsewhere. Multiplying by that row keeps only the log-probability of the right digit.
• The log-softmax. The program computes log p directly as z - log Σ e^z per row, with z already shifted by the max. That is the log of softmax, without the division. The loss is minus the mean of the picked values.
• The permutes. reduce pops leading axes, and the classes are the trailing axis, so each per-row reduction transposes first.
• The stage. Without it, fusion would recompute each logit (a 784-term dot product) at every use. Because the gradient of a staged value is staged too, the backward pass gets the same benefit.
• The update. It reads %w and stores into %w. That's allowed: a store's value is computed from the contents before it. Your renderer must keep it that way when the store becomes a loop (lesson 6.6).

Fifty steps, written out

Step 1 is the same with s1_ names, batch 1 (shrink %xs 78400), and the weights step 0 left: load %s0_w, stored into %s0_w. Week 6 has no loops around whole tensors, so all 50 steps are written out, about 2000 lines. Every step is the same shape of graph, so most of the work your compiler does on step 0 it does again 49 times. It needs to be fast enough to render the whole program in a reasonable time, and the test gives it 120 seconds.

After step 49, the program runs the 1000 test images through the final weights. A digit counts as right when the logit of its true label is the largest. The count goes into slot 4. The reference gets 864, and the test needs at least 850.