Kahibaro
Discord Login Register

13.5 Batch Inference and Latency Considerations

What batch inference means in practice

Batch inference is running a trained model on multiple inputs at once. Instead of calling the model for a single example, you stack many examples into a batch tensor and run one forward pass. This is usually faster overall because it reduces Python overhead and uses the hardware more efficiently, especially on GPUs.

Latency is the time it takes to return a prediction. Throughput is how many predictions you can produce per second. Batching tends to increase throughput, but it can increase per request latency if you wait to collect a batch. Deployment work is often about balancing these two goals.

Latency versus throughput tradeoffs

If you make the batch size larger, the model often runs more efficiently, so throughput improves. However, the system may need to wait until enough requests arrive to form a batch, which adds queueing delay. For interactive use, you usually care about tail latency, such as the 95th or 99th percentile, not just the average. For offline jobs, you usually care about throughput and cost.

A practical mental model is that total latency is compute time plus overhead plus queueing time. Batching mainly changes compute time and queueing time, and the best batch size depends on your hardware, model, and traffic pattern.

Setting up PyTorch for fast inference

For inference you should disable gradient tracking and switch the model to evaluation mode. This avoids autograd overhead and ensures layers like dropout and batch normalization behave correctly.

Always use model.eval() for inference and wrap forward passes in torch.inference_mode() or torch.no_grad(). Forgetting this can increase latency and memory usage and can change predictions.

In PyTorch, torch.inference_mode() is usually preferred for pure inference because it can be faster and use less memory than torch.no_grad() in many cases.

Choosing a batch size that makes sense

There is no universal batch size. You pick one that fits in memory and meets your latency target. A good approach is to benchmark a few candidate batch sizes, for example 1, 2, 4, 8, 16, 32, and measure both throughput and latency. On GPUs, batch size 1 can be surprisingly inefficient. On CPUs, batching helps too, but the optimal point can be smaller.

When memory is a constraint, remember that activation memory at inference is usually smaller than training, but it still grows with batch size. Also note that some models, especially transformer based models, have memory use that grows with sequence length, so a larger batch can push you over the limit quickly.

Dynamic batching and micro batching

If you serve online requests, you may implement dynamic batching. The server collects requests for a short time window and merges them into one batch. This can improve throughput while keeping latency under control. A common pattern is to use a maximum batch size and a maximum waiting time. If either threshold is reached, the batch is executed.

Micro batching is a related idea where a large batch is split into smaller chunks to fit in memory. You run the model multiple times and concatenate outputs. This increases overhead slightly but prevents out of memory errors.

If you implement micro batching, ensure that your output ordering matches your input ordering. Mixing up indices is a common, hard to detect production bug.

CPU versus GPU inference considerations

GPUs excel at parallel work, so batching usually helps more on GPU than on CPU. However, moving data from CPU to GPU adds overhead. If each request is tiny, the transfer cost can dominate. This is why batch size 1 GPU inference can be slower than CPU for small models or small inputs.

If you do use a GPU, keep tensors on the GPU as much as possible. Repeated .to(device) calls inside the per request path add overhead. It is often better to move a whole batch once and keep the model on the same device.

Also remember to synchronize carefully when measuring performance. GPU operations are often asynchronous, so naive timing can be wrong unless you synchronize around the timed region.

Input pipeline overhead and DataLoader choices

For offline batch inference, input parsing and preprocessing can become the bottleneck. If the model is fast, CPU preprocessing, image decoding, tokenization, and collation can dominate total time. In those cases, increasing batch size alone will not fix throughput.

For offline inference with a DataLoader, you may reuse performance ideas you already use for training, such as using multiple workers and pinned memory when transferring to CUDA. The key difference is that you typically do not need shuffling and you often want deterministic ordering so you can align outputs back to original inputs.

Avoiding unnecessary work inside the hot path

Inference code should minimize Python overhead. Repeatedly creating small tensors, repeatedly allocating lists, and converting between NumPy and PyTorch can add noticeable latency. Prefer vectorized operations on tensors and batch preprocessing where feasible.

You should also avoid calling softmax if you only need the argmax class, because argmax of logits matches argmax of softmax. Similarly, if you only need a top class and not probabilities, skipping probability normalization reduces compute.

For classification, argmax(logits) equals argmax(softmax(logits)). Do not compute softmax unless you actually need probabilities.

Measuring latency correctly

You should separate warmup from measurement. The first few runs can be slower due to caching, kernel compilation, or initialization overhead. For fair comparisons, run several warmup iterations and then time many iterations.

Measure end to end latency if you are making deployment decisions. That means including preprocessing, device transfer, model forward pass, and postprocessing. If you only time the model forward, you might optimize the wrong component.

When measuring on GPU, use synchronization around timing so you measure actual completion time rather than queued work.

Common deployment patterns and what they imply

In offline inference, you typically maximize throughput. You can use larger batch sizes, fewer constraints on latency, and you can schedule jobs to keep the GPU saturated.

In online inference, you often use smaller batches or dynamic batching. You focus on predictable latency, memory stability, and fast startup. You may also run multiple model replicas to reduce queueing delay.

In edge or mobile scenarios, CPU inference, quantization, and smaller models are common, and batch size is often 1 because requests are single and interactive. In that context, you optimize per sample latency rather than throughput.

Practical rules of thumb

Optimize in this order: correctness first, then measure end to end, then remove the biggest bottleneck. Guessing usually wastes time.

A good default starting point is to implement inference with model.eval() and torch.inference_mode(), support a configurable batch size, and benchmark a small grid of batch sizes on your target hardware while tracking both throughput and latency percentiles.

Views: 58

Comments

Please login to add a comment.

Don't have an account? Register now!