How to use a Predictor to predict on batches?

[hmf]

I am experimenting with Autoencoders on the MNIST dataset. I can successfully train and test the whole model by evaluating the L2 loss on the image reconstruction. For the prediction I use a predictor so:

      val predictor: Predictor[NDArray, NDArray] = model.newPredictor(translator)

and my translator is:

    val translator = new Translator[NDArray, NDArray]:

      override def processInput(ctx: TranslatorContext, input: NDArray): NDList =
        // Reshape to correct input size
        NDList(input.reshape(Shape(1, Mnist.IMAGE_HEIGHT * Mnist.IMAGE_WIDTH)))

      override def processOutput(ctx: TranslatorContext, output: NDList): NDArray = 
        // Reshape to correct output size
        output.head.reshape(Shape(Mnist.IMAGE_HEIGHT * Mnist.IMAGE_WIDTH))
      
      override def getBatchifier(): Batchifier =
        // The Batchifier describes how to combine a batch together
        // Stacking, the most common batchifier, takes N [X1, X2, ...] arrays to a single [N, X1, X2, ...] array
        Batchifier.STACK

The batchifier generates batches so that each NDarrayhas the shape (256, 1, 28, 28). So when I want to predict the output I do the following:

          // Get the batch
          val images = batch.getData.head
          val labels = batch.getLabels.head
          // get the first element of the batch
          val image = images.get(0).squeeze(0)
          val label = labels.get(0).squeeze(0)

          // predict (reconstruction): label and image are the same
          val predicted = predictor.predict(label)
          val deltaLoss = loss.evaluate(NDList(image), NDList(predicted)).getFloat()

In other words I take a batch of 256 images, sample the first one of Shape(1,28,28). The translator turns that into an input of Shape(784) and gets an output of Shape(784). The translator then converts that back to Shape(28, 28). This works correctly.

Now what I would like to do is take the full Shape(256, 1, 28, 28) use that as input of Shape(256, 784) to the predictor and get an output of Shape(256, 784).

So the first issue is that the code of the translator has the getBatchifier, which seems to suggest that the predictor is already prepared to use batches. However, I don't see how this can be used. Is their an example of using the Predictor with batches?

I have also attempted to use a translator that does the following: it takes the Shape(256, 1, 28, 28) input and converts that to Shape(256, 784). I now expect an output of Shape(256, 784). However, what I get from Predictor.predict is a Shape(784). What am I doing wrong?

Finally, how does one handle the last batch that may not have the required length of 256?

TIA

[zachgk]

To use the predictor with batches, just call predictor.batchPredict(...) instead. In this case, you would want to use a translator that translates single images at a time and the StackBatchifier. It batches any size so no modifications are necessary for the last partial batch.

[zachgk]

The way that it is expected to work is:

• Input to batchPredict of List<size 1x28x28 NDArray images>
• Translator preprocesses each image from Shape(1, 28, 28) to Shape(784) so you have a List<Shape(784) NDArray images>
• The stack batchifier stacks the images together to form a Nx748 NDArray with batch size N
• Run the model to produce some output as a joined NDArray NxM with output size M (I think your M is 784 for an autoencoder, but would be 10 for a classifier).
• The stack batchifier unstacks the images back into a List<Shape(M) NDArray
• The translator postprocesses each element in the List<> to form the output types
• The outputs are returned

---

One option you can take if you don't like this process is skipping the batchifier. If your translator has null as the batchifier, then you can pass in a batch Nx784 into the usual predictor.predict() call. Your translators should then translate batches instead of single images.

[hmf]

@zachgk Apologies for removing my previous question, but I realized that the stacking was done inside the predictor and was thinking about this. Ok, I finally get it. I will code up both techniques because it will serve as documentation. The "null" technique is useful because I am using a RandomAccessDataset to load the data set. Very useful.

Appreciate the help.
Thank you.