We had settled the hard modelling questions before this run and were left with a plainer one that turned out to matter just as much: could we train the multiclass segmenter fast enough to iterate on it. The dataset was 15 thousand synthetic log instances, the target was three classes, and the schedule was fifty epochs. None of that was in question here. The problem was that the first version of the training loop treated the 15 thousand instances more or less one at a time, and on images this large that is not a training run, it is a vigil. This piece is only about the speed. Not the loss, which was decided elsewhere, and not the choice to move from binary masks to a three-class target, which was its own decision with its own reasons. Just the wall-clock, and the handful of levers that pulled it down to about ten hours.
Why one image at a time was the whole problem
The images are the reason the naive loop was slow, and they are worth describing precisely because they dictate everything that follows. A synthetic log instance is tall and narrow and, crucially, not a fixed size. Widths ran from 3,200 to 12,800 pixels and heights from 480 to 640, so no two batches held images of the same shape. That variability is fatal to the ordinary way you accelerate training. The standard speedup is to stack many samples into one batch tensor and let the GPU work on all of them in a single pass, but you cannot stack tensors of different widths and heights into one rectangular array. The path of least resistance is to give up on batching and feed the network one image per step, and that is exactly where the binary run had lived: batch size one, pinned there not by choice but by the fact that variable image sizes left no obvious way to group them.
Batch size one is the slowest place a training loop can be. Every one of the 15 thousand instances becomes its own optimiser step, so a single epoch is 15 thousand steps, and fifty epochs is three quarters of a million of them. Each step carries fixed overhead that has nothing to do with the arithmetic of the forward and backward pass, the launch of the kernels, the synchronisation, the gradient update, and at batch one you pay that overhead 15 thousand times per epoch instead of amortising it across a group. The GPU spends much of its time waiting rather than computing. We had a concrete reference point for how this felt: the earlier binary run, on two thousand instances at batch one for the same fifty epochs, took 110 minutes. Scale that per-instance cost up to 15 thousand instances and the multiclass run was heading somewhere we could not plan a week around.
The lever that mattered was the batch, and the collate_fn is what unlocked it
The single relationship that governs this whole story is simple. Total wall-clock is roughly the number of optimiser steps times the cost per step, and the number of steps per epoch is the instance count divided by the batch size. Every fixed-overhead cost you pay per step gets divided by the batch. Take the batch from one to 16 and you have cut the number of steps per epoch by a factor of 16, from 15 thousand down to under a thousand, and you have handed the GPU 16 images to work on in parallel where before it had one. That is the entire mechanism. Everything else we did was in service of making a batch of 16 possible at all on images this awkward.
Making it possible came down to a custom collate_fn. A dataloader's collate function is the step that takes a list of individual samples and assembles them into one batched tensor, and the default one assumes every sample is the same shape. Ours could not assume that, so we wrote a collate_fn that took a group of 16 variable-dimension images and padded them to a common shape before stacking, recording enough to keep the padding from being mistaken for signal downstream. With that in place the dataloader could actually return batches of 16 instead of surrendering to batch one, and the throughput relationship above stopped being theoretical.
The three levers that kept a batch of sixteen inside the memory budget
A batch of 16 large images is a lot of memory, and the reason batch one had been the default was never really the collate logic alone, it was that bigger batches did not fit. So the collate_fn came with company. Mixed precision let us hold activations in half the space by using 16-bit floats where full precision was not needed, which roughly halved the activation footprint and, on the right hardware, sped up the math as a bonus rather than a cost. Gradient checkpointing bought more headroom by refusing to keep every intermediate activation from the forward pass in memory, instead recomputing the cheap ones during the backward pass, a deliberate trade of a little extra compute for a lot less memory. And the dataloader itself was tuned to prepare the next batch while the current one trained, so the card was not left idle waiting on the CPU to pad and stack the next 16 images.
None of these three changed the model's output. Mixed precision, checkpointing, and prefetching are all invisible to the loss and the weights: the network sees the same images, computes the same gradients to within numerical tolerance, and arrives at the same place. What they changed was whether a batch of 16 fit in the memory budget, and that is the only thing that let the batch lever move. Take any one of them away and the batch that fits shrinks back toward one, and the wall-clock climbs back up the curve the instrument above traces.
What the clock actually read
With the collate_fn, batch 16, mixed precision, checkpointing, and a fed dataloader all in place, the multiclass run over 15 thousand instances for fifty epochs took 550 minutes, a little over nine hours, which we rounded to about ten hours when we talked about it as an overnight job. The comparison that made the point internally was against the binary anchor: 110 minutes for two thousand instances at batch one. The multiclass set was seven and a half times larger and the run took five times as long, not the seven and a half or more you would expect if per-instance cost had held constant. The batch had absorbed the difference. We used an 80/20 train-validation split throughout, so the epoch was training on twelve thousand of the 15 thousand instances, and that is the number the throughput math above is really dividing.
The reason ten hours mattered was not the ten hours. It was that ten hours is one night. A run that finishes overnight is a run you can start at the end of a day, read in the morning, and change something about before the next evening, which turns a training loop into something you can iterate against on a daily cadence. Two days per run does not give you that. It gives you a couple of attempts a week and a strong incentive to guess rather than measure. Pulling the wall-clock under a working night was the difference between a model we tuned by experiment and one we would have tuned by superstition.
Limitations
The specific 550-minute figure is one number from one engagement on one machine, and it should be read as evidence that these levers move wall-clock by roughly this much on a workload like this, not as a benchmark to port. The projection between the two measured points in the exhibit is a step-count model calibrated to reproduce the measured floor, not a per-batch timing; the real curve would bend where memory bandwidth, kernel efficiency, or dataloader throughput become the binding constraint rather than step count, and we did not sweep the intermediate batch sizes to map that bend. Larger batches also change the effective gradient noise, so beyond a point raising the batch is no longer free with respect to what the model learns, even though it stays cheap with respect to the clock; on this run 16 was the ceiling memory allowed and we did not need to push against the learning question. Finally, everything here rests on the synthetic image dimensions we describe, 3,200 to 12,800 pixels wide and 480 to 640 tall. Field scans with a wider spread of sizes would put more pressure on the padding in the collate_fn, and the memory headroom that made batch 16 fit is not guaranteed to survive a heavier tail of very wide images.
