How can I use a model trained with batches to make predictions with single sample?

Question

I'm training a PyTorch model with batches of 128 images, and after going through multiple convolutions, they're flattened (with .flatten) before being passed to a linear layer. The input size of the linear layer accounts for the batch size, so the dimensions don't line up when I pass a single image to the trained model. How is this typically dealt with?

Answer 1

Your scenario is extremely common. You generally want to train in batches and then predict on individual samples.

First, it's crucial to understand that training and using a model are separate endeavors. While you might be able to reuse code, it is in no way a requirement.

I like how PyTorch lightning handles it. Essentially, they define two methods in their interfaces, training_step for training and predict_step for inference. As suggested, they deal with both aspects. And what you could always do is use the predict_step as part of your training_step to avoid duplicate code.

For your use case, I would make sure that your prediction method is flexible when it comes to batch size (so when you predict, you pass a batch of 1 item) or I would design two separate methods in case the logic between the two varies enough.

Answer 2

I suggest flattening with .flatten(start_dim=1) instead (keep the batch dimension 128 unaffected, flatten information for each image only). We should not have the last linear layer input size dependent on the batch dimension. Check out this implementation of ResNet in PyTorch, see line 112 of ResNet PyTorch Implemnetation. Our model should be compatible with any input batch size.

Think about linear regression with design matrix $X \in \mathbb{R}^{n \times p}$, with $n$ samples each with $p$ dimension. We want to train a weight vector $\beta \in \mathbb{R}^p$ that should not depend on $n$.

Using .flatten(start_dim=1) makes the input to the linear layer a $n$ by $p$ matrix, where $n$ is the batch size 128 here, and $p$ is determined by the input dimension and convolution settings.