Train a classifier on specific labels of MNIST dataset with TensorFlow

I would like to train a classifier on MNIST dataset but with limited labels. For eg. I would like to train a classifier only on labels [1, 4, 5, 6, 8, 9] from all the labels [0-9]. I am getting the following error:

res = tf.nn.sparse_softmax_cross_entropy_with_logits(
Node: 'sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits'
Received a label value of 9 which is outside the valid range of [0, 6). Label values: 9 5 9 9 6 1 4 4 6 6 9 4 9 1 8 5 9 5 4 8 9 9 1 8 6 4 4 9 9 4 4 8 8 6 6 5 9 4 1 5 5 6 4 1 1 8 9 6 8 5 6 1 6 6 4 6 1 4 4 4 1 1 1 6 9 8 8 8 5 1 8 8 6 6 5 1 1 5 1 6 9 8 1 8 4 6 4 9 8 1 6 5 5 9 1 6 8 1 5 5 6 9 1 9 9 6 4 6 6 4 8 6 6 4 5 4 4 5 8 1 8 6 1 5 4 5 8 1

Here is approach I have used:

import tensorflow_datasets as tfds
import tensorflow as tf
val_split = 20  # percent of training data
(ds_test, ds_valid, ds_train), ds_info = tfds.load(
    'mnist,
    split=['test', f'train[0%:{val_split}%]', f'train[{val_split}%:]'],
    as_supervised=True,
    with_info=True
)

The ds_train dataset object has the following samples per label <br>
{0: 4705, 1: 5433, 2: 4772, 3: 4936, 4: 4681, 5: 4333, 6: 4728, 7: 4966, 8: 4703, 9: 4743}.

After this I filter the dataset using filter() as follows:

known_classes = [1, 4, 5, 6, 8, 9]
kc = tf.constant(known_classes, dtype=tf.int64)
def predicate(image, label):
    isallowed = tf.equal(kc, label)
    reduced = tf.reduce_sum(tf.cast(isallowed, tf.int64))
    return tf.greater(reduced, tf.constant(0, dtype=tf.int64))
ds_test = ds_test.filter(predicate)
ds_valid = ds_valid.filter(predicate)
ds_train = ds_train.filter(predicate)

The updated samples per label post filter fo ds_train is <br>
{0: 0, 1: 5433, 2: 0, 3: 0, 4: 4681, 5: 4333, 6: 4728, 7: 0, 8: 4703, 9: 4743}.

Next steps are normalizing the image and preparing the dataset objects for training.

def normalize_img(image, label):
    """Normalizes images: `uint8` -> `float32`."""
    return tf.cast(image, tf.float32) / 255., label
ds_train = ds_train.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds_train = ds_train.shuffle(ds_info.splits['train'].num_examples)
ds_train = ds_train.batch(128, drop_remainder=True)
ds_train = ds_train.prefetch(tf.data.AUTOTUNE)
ds_test = ds_test.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds_test = ds_test.batch(128, drop_remainder=True)
ds_test = ds_test.prefetch(tf.data.AUTOTUNE)
ds_valid = ds_valid.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds_valid = ds_valid.batch(128, drop_remainder=True)
ds_valid = ds_valid.prefetch(tf.data.AUTOTUNE)

Thereafter, I create a simple model as follows and then proceed with training

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dense(len(known_classes))  # known_classes from above
])
model.compile(
    optimizer=tf.keras.optimizers.Adam(0.001),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
)
history = model.fit(
    ds_train,
    epochs=5,
    validation_data=ds_valid,
)

I am new to Tensorflow and any help is appreciated!!

You have to change the training labels too, mapping them to 0-5 range instead of [1, 4, 5, 6, 8, 9] because the sparse_categorical_crossentropy expects them to be in 0-5 since you removed some of them.

You can do this when normalizing the images, for example:

def make_fn(y):
    def fn(z):
        return tf.equal(z, tf.constant(y, dtype=tf.int64))
    return fn
def make_case(i):
    def case():
        return tf.constant(i, dtype=tf.int64)
    return case
label_cases = [(make_fn(y), make_case(i)) for i, y in enumerate(known_classes)]
def normalize_img_and_label(image, label):
    """Normalizes images and labels: `uint8` -> `float32`."""
    label = tf.case([(fn(label), case) for fn, case in label_cases], 
                    exclusive=False, default=lambda: label)
    return tf.cast(image, tf.float32) / 255., label
ds_train = ds_train.map(normalize_img_and_label, num_parallel_calls=tf.data.AUTOTUNE)
...

Indeed, you do this also on ds_valid and ds_test

Explanation: we want to implement the following mapping of the remaining class labels, i.e. [1 -> 0, 4 -> 1, 5 -> 2, 6 -> 3, 8 -> 4, 9 -> 5], because the loss function expects them to be in 0-5 range. Because of functions that operates on tf.data.Dataset objects (e.g. normalize_img()) work on symbolic tensors we can't implement the mapping directly using a dict with integers key-value pairs (as my previous solution.)

Within the normalize_img_and_label function both image and label are symbolic, meaning that they are like placeholders with no numpy value associated to them: such value will be available when actually consuming the dataset object. So we must use tf primitives able to handle symbolic tensors. Here I make use of a tf.case to implement a switch-case statement over symbolic tensors, of the kind:

if label == 1:
  return 0
elif label == 4:
  return 1:
...
elif label == 9:
  return 5

To do this you define a list of function pairs (i.e. label_cases): one pair for each label.

• The make_fn creates a closure that when called, e.g. fn(label), return the value of the boolean conditional tf.equal(z, y) - where y is one value of known_classes and z is the provided label.
• The make_case yields another closure that returns a constant tensor, corresponding to the return value in the switch-case, i.e. the new mapped label.

Now, inside normalize_img_and_label we have to apply the tf.case and so we loop over the pairs of functions, yielding a new pair (Boolean, tensor) since we have to pass the label tensor to the tf.equal closure (i.e. fn.) For example assuming label = 4, the list [(fn(label), case) for fn, case in label_cases] is actually: [(False, 0), (True, 1), (False, 2), ..., (False, 5)]. Therefore the tf.case would yield a 1 because of the true condition.