Training QuickDraw with PyTorch

1. From stroke to image

Quick, Draw! stores drawings as stroke sequences: lists of x and y coordinates interrupted by pen lifts. For a CNN these must become an image of fixed size. This step decides more about recognition quality than the choice of architecture, and it is often treated too casually.

The crucial part is normalisation. The bounding box is computed across every point of the drawing, the drawing is shifted to the centre and scaled to fit with an eight per cent margin. The scale factor is the same for both axes — otherwise a house drawn wide would be squashed into a square one and would no longer be the same object:

span_x = max(max_x - min_x, 1e-6)
span_y = max(max_y - min_y, 1e-6)
pad    = max(1.0, img_size * 0.08)
avail  = max(2.0, (img_size - 1) - 2.0 * pad)
scale  = min(avail / span_x, avail / span_y)   # one scale for both axes

The result: where on the canvas someone draws, and how large, no longer matters. Only shape counts. The segments are then drawn with a line width of two pixels, setting every point along a straight line between consecutive vertices.

It is essential that training and inference use exactly the same code. In DrawLa the rasteriser therefore lives in a single module (server/preprocess.py) which the training scripts import. If the two paths diverge by even one pixel of margin, accuracy drops noticeably without anything looking wrong during training — a bug that can hide for a long time.

2. Why 96 × 96

The image size is a compromise. Smaller means faster and less memory, but fine distinctions disappear: at 28 × 28 — the MNIST size many people reach for — the spokes of a bicycle merge into a grey smudge. Larger means more compute per prediction, which is felt directly when recognising eight times per second.

96 × 96 has proved a workable middle ground: enough resolution for details such as windows, wheels or legs, and small enough that a prediction stays in the single-digit millisecond range on an ordinary server CPU. A single greyscale channel suffices, since colour carries no information in line drawings.

3. The data: one NPZ file per category

The dataset covers 345 categories. At 5,000 drawings per class that comes to roughly 1.7 million images. Rasterising happens once, up front, not on every training step — putting it in the dataloader would be a reliable way to starve the GPU while the CPU draws lines.

Each category is stored as a compressed NPZ file holding a uint8 array of shape (N, 96, 96). One byte per pixel instead of four: the whole set stays manageable, and the dataloader converts to float only on access.

def __getitem__(self, i):
    img_u8, label_idx = self._samples[i]
    x = torch.from_numpy(img_u8.astype(np.float32) / 255.0).unsqueeze(0)
    if self.augment:
        x = self._augment_image(x)
    return x, label_idx

The split into training and validation is 90 to 10 within each category. A global random split would over- or under-represent individual classes in validation whenever category sizes differ.

4. Augmentation: what makes sketches real

People draw crooked, skewed and at varying sizes. That is exactly what augmentation has to reproduce. Each image gets a random rotation of ±15°, a scale between 0.9 and 1.1, an independent stretch of both axes in the same range, and a translation of up to 16 per cent of the image edge.

Two details are specific to line drawings. First, line-width jitter: with 20 per cent probability each, the line is thickened or thinned by one pixel, implemented as max pooling on the image and on its negative respectively. This mimics the fact that fingers on a touchscreen produce thicker strokes than a mouse pointer.

width_jitter = int(self.rng.choice([-1, 0, 1], p=[0.2, 0.6, 0.2]))
if width_jitter > 0:                      # thicker: lines grow
    image = F.max_pool2d(image.unsqueeze(0), 3, stride=1, padding=1).squeeze(0)
elif width_jitter < 0:                     # thinner: dilate the negative
    inv = 1.0 - image
    inv = F.max_pool2d(inv.unsqueeze(0), 3, stride=1, padding=1).squeeze(0)
    image = 1.0 - inv

Second, restraint with noise: only in a quarter of cases, and then with a standard deviation of at most 0.02. Line drawings are empty almost everywhere; heavy noise fills that emptiness with structure that never occurs in real input and makes results worse. Mirroring is deliberately absent — it would turn a "b" into a "d" and destroy direction-dependent objects.

5. The model: 2.0 million parameters

The architecture is a scaled-down ResNet. A stem immediately halves the resolution, followed by four stages of two residual blocks each; the channel count rises through 32, 64, 128 to 192 while resolution falls from 96 to 6 pixels. Global average pooling, dropout and a linear layer onto 345 classes finish it off.

Stem  : Conv(1→32, 3×3, stride 2) + BatchNorm + SiLU
Stage1: 2 × ResidualBlock, 32 channels          48×48
Stage2: 2 × ResidualBlock, 64 channels          24×24
Stage3: 2 × ResidualBlock, 128 channels         12×12
Stage4: 2 × ResidualBlock, 192 channels          6×6
Head  : GlobalAvgPool → Dropout(0.2) → Linear(192→345)

That yields 2,004,345 parameters, exactly 8.0 MB in fp32. For comparison: a ResNet-50 has roughly twelve times as many and would be considerably oversized here — line drawings have no textures, no lighting and no background, only shape.

Global average pooling instead of a large fully connected layer is especially effective here: it reduces 6 × 6 × 192 values to 192 and thus saves the lion's share of the parameters a classical classifier would spend at this point. SiLU instead of ReLU gives slightly smoother gradients, and batch normalisation keeps training stable even at a learning rate of 10⁻³.

6. Training, and what comes out of it

Training uses AdamW at a learning rate of 10⁻³ with weight decay 10⁻⁴, over 40 epochs with a batch size of 256. The learning rate follows a cosine schedule down to 10⁻⁵. The loss is cross-entropy with label smoothing 0.1 — sensible because the dataset itself is mislabelled in places: many drawings are abandoned or simply failed, and a model trained to full confidence on such labels becomes overconfident.

criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler = CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-5)
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

The best result was 76.3 per cent top-1 accuracy on the validation split, reached in epoch 39 of 40. That number seems modest until you put it in context: with 345 equally likely classes, guessing lands at 0.29 per cent. The model is therefore about 260 times better than chance, and much of the remainder is not a weakness of the network but ambiguity in the source material. A quick scribble of "snail" and one of "seashell" are sometimes the same picture.

In the game the top-5 list is what counts anyway, and there the hit rate is considerably higher. Only the best checkpoint is saved, together with labels and image size — so the export cannot accidentally work with a different class ordering later.

7. Exporting to ONNX

For production the model is exported to ONNX with opset 18 and a dynamic batch axis. That decouples inference from PyTorch: no Torch runs on the server, which shrinks the container image considerably and shortens cold starts.

torch.onnx.export(
    model, torch.zeros(1, 1, 96, 96), "models/quickdraw_345cls.onnx",
    opset_version=18,
    input_names=["image"], output_names=["logits"],
    dynamic_axes={"image": {0: "batch"}, "logits": {0: "batch"}},
)
onnx.checker.check_model(onnx.load(out_path))   # never ship unchecked

The labels are written as JSON next to the model in the same step. Class order and weights belong together; maintaining them separately is a dependable source of bugs that surface as "the model suddenly gets everything wrong".

8. Inference during a live game

In production, ONNX Runtime runs on the CPU with two threads per request and graph optimisation fully enabled. A GPU would be uneconomical for a model this size — the cost of moving data back and forth outweighs the computation.

The client sends the current stroke state every 120 milliseconds, a good eight times per second. The logits pass through a softmax that subtracts the maximum first — otherwise exp() can overflow:

logits_shifted = logits - logits.max()      # numerically stable
probs = np.exp(logits_shifted)
probs /= probs.sum()

One final trick makes the difference to how the game feels: exponential smoothing across consecutive predictions with α = 0.6. Without it the display flickers between competing words with every new stroke, which feels like the model is guessing. With it, the recognition appears to form a hypothesis and stick to it until enough evidence accumulates against it. The smoothing is purely presentational — it changes nothing about accuracy and a great deal about perceived behaviour.

What we would do differently

In hindsight, compute spent on data quality would have paid off better than compute spent on architecture. The Quick, Draw! dataset contains a noticeable share of abandoned and mislabelled drawings; filtering out the most obvious cases would probably have helped more than any additional residual stage.

Weighting by confusability would also be worthwhile. A few class pairs are practically inseparable, and the model wastes capacity trying to separate them anyway instead of becoming more confident on the unambiguous categories.

How this image-based approach compares to the sequence-based alternative is covered in RNN versus CNN. The data is explained in the Quick, Draw! dataset, and why such a small model is a good candidate for on-device inference is discussed in Mobile & Edge AI.