how to get labels when using model.predict()-CodePudding

In my project, I have a number of cases where I have a Dataset instance and I need to get predictions from some model on every item in the dataset.

The model.predict() API is optimized perfectly for this, as shown in the documentation. However, there seems to be one major catch. I also happen to need the labels to compare with the predicted values, i.e. the dataset contains x,y pairs, and I'd like to end up with (y_predicted, y) pairs after the prediction is complete. This does not seem to be possible with the predict() API though, and I can't think of a clean way to 'split' the dataset so that the x's are fed into the model and the y's are retained to be joined back up with the predicted y's.

EDIT: I know it's quite simple to do by iterating over the dataset manually and calling the model directly, e.g.

for x, y in dataset:
    y_pred = model(x)
    result.append((y, y_pred))

However, this seems like it will be a fair bit slower than using the inbuilt predict() as Tensorflow won't be able to multi-thread/optimize the input pipeline.

Does anyone have a good way to accomplish this?

CodePudding user response：

Given the concerns you mentioned, it may be best to overwrite predict to suit your needs. You don't actually need to overwrite that function though, instead only predict_step which is called by that function. Just use this class instead of Model:

class MyModel(tf.keras.Model):
    def predict_step(self, data):
        x, y = data
        return self(x, training=False), y

If your model is currently Sequential, inherit from that instead. Basically the only change I made from the default implementation is to add , y to the model call result. Note that this also makes some assumptions, such that your dataset consists of (input, label) batch pairs. You may need to adapt it slightly to your needs. Here is a minimal example:

import tensorflow as tf
import numpy as np


(imgs, lbls), (te_imgs, te_lbls) = tf.keras.datasets.mnist.load_data()

imgs = imgs.astype(np.float32).reshape((-1, 784)) / 255.
te_imgs = te_imgs.astype(np.float32).reshape((-1, 784)) / 255.

lbls = lbls.astype(np.int32)
te_lbls = te_lbls.astype(np.int32)

tr_data = tf.data.Dataset.from_tensor_slices((imgs, lbls)).shuffle(60000).batch(128)
te_data = tf.data.Dataset.from_tensor_slices((te_imgs, te_lbls)).batch(128)


class MyModel(tf.keras.Model):
    def predict_step(self, data):
        x, y = data
        return self(x, training=False), y


inp = tf.keras.Input((784,))

logits = tf.keras.layers.Dense(10)(inp)

model = MyModel(inp, logits)


opt = tf.keras.optimizers.Adam()
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

model.compile(loss=loss, optimizer=opt)

something = model.predict(te_data)

print(something[0].shape, something[1].shape)

This shows ((10000, 10), (10000,)) -- predict now returns a tuple of outputs, labels (this can be confirmed by inspecting the returned labels and comparing to the images in the test set).