Attaching class labels to a Keras model

Viewed 4554

I am using the Keras Sequential model to train a number of multiclass classifiers.

On evaluation, Keras outputs a vector of confidences and I can infer the correct class id from that using argmax. I can then use a lookup table to receive the actual class label (e.g. a string).

So far the solution is to load the trained model, and then to load a lookup table separately. Since I have quite a number of classifiers I would prefer to keep both structures in one file.

So what I am looking for is a way to integrate the actual label lookup vector into the Keras model. That would allow me to have a single classifier file that is capable of taking some input data and returning the correct class label for that data.

One way to solve this would be to store both the model and the lookup table in a tuple and write that tuple into a pickle, but this doesn't seem very elegant.

2 Answers

It is possible to save a "list" of labels in keras model directly. You need to use lambda. You just replace the output of lambda with a string tensor containing your labels. Here is a dummy example of how one can perform an "injection" of labels

# assume we get labels as list
labels = ["cat","dog","horse","tomato"]
# here we start building our model with input image 299x299 and one output layer
xx = Input(shape=(299,299,3))
flat = Flatten()(xx)
output = Dense(shape=(4))(flat)
# here we perform injection of labels
tf_labels = tf.constant([labels],dtype="string")
# adding ? dimension to tf tensor
tf_labels = tf.tile(labels,[tf.shape(xx)[0],1])
output_labels = Lambda(lambda x: tf_labels,name="label_injection")(xx)
#and finaly creating a model
model=tf.keras.Model(xx,[output,output_labels])

This model now stores the labels and returns them as well. All this mess with tf.tile is necessary because a keras layer of shape (N) is actually a tf tensor of shape (?,N) and we add this ? dimension to the label tensor.

Related