Multioutput-Multiclass Classification in Custom Scratch Training in TF.Keras

Question

I want to train a multi-out and multi-class classification model from scratch (using custom fit()). And I want some advice. For the sake of learning opportunity, here I'm demonstrating the whole scenario in more detail. Hope it may come helpful to anyone.

Data Set and Goal

I'm using data from here; It's a Bengali handwritten character recognition challenge, each of the samples has 3 mutually related output along with multiple classes of each. Please see the figure below:

In the above figure, as you can see, the ক্ট্রো is consist of 3 component (ক্ট , ো , ‍‍্র), namely Grapheme Root, Vowel Diactrics and Consonant Diacritics respectively and together they're called Grapheme. Again the Grapheme Root also has 168 different categories and also same as others (11 and 7). The added complexity results in ~13,000 different grapheme variations (compared to English’s 250 graphemic units).

The goal is to classify the Components of the Grapheme in each image.

Initial Approach (and no issue with it)

I implemented a training pipeline over here, where it's demonstrated using old keras (not tf.keras) with its a convenient feature such as model.compile, callbacks etc. I defined a custom data generator and defined a model architecture something like below.

input_tensor = Input(input_dim)
curr_output = base_model(input_tensor)

oputput1 = Dense(168,  activation='softmax', name='gra') (curr_output)
oputput2 = Dense(11,   activation='softmax', name='vow') (curr_output)
oputput3 = Dense(7,    activation='softmax', name='cons') (curr_output)
output_tensor = [oputput1, oputput2, oputput3]
    
model = Model(input_tensor, output_tensor)

And compile the model as follows:

model.compile(

        optimizer = Adam(learning_rate=0.001), 

        loss = {'gra' : 'categorical_crossentropy', 
                'vow' : 'categorical_crossentropy', 
                'cons': 'categorical_crossentropy'},

        loss_weights = {'gra' : 1.0,
                        'vow' : 1.0,
                        'cons': 1.0},

        metrics={'gra' : 'accuracy', 
                 'vow' : 'accuracy', 
                 'cons': 'accuracy'}
    )

As you can see I can cleary control each of the outputs with specific loss, loss_weights, and accuracy. And using the .fit() method, it's feasible to use any callbacks function for the model.

New Approach (and some issue with it)

Now, I want to re-implement it with the new feature of tf.keras. Such as model subclassing and custom fit training. However, no change in the data loader. The model is defined as follows:

    def __init__(self, dim):
        super(Net, self).__init__()
        self.efnet  = EfficientNetB0(input_shape=dim,
                                     include_top = False, 
                                     weights = 'imagenet')
        self.gap     = KL.GlobalAveragePooling2D()
        self.output1 = KL.Dense(168,  activation='softmax', name='gra')
        self.output2 = KL.Dense(11,   activation='softmax', name='vow') 
        self.output3 = KL.Dense(7,    activation='softmax', name='cons') 
    
    def call(self, inputs, training=False):
        x     = self.efnet(inputs)
        x     = self.gap(x)
        y_gra = self.output1(x)
        y_vow = self.output2(x)
        y_con = self.output3(x)
        return [y_gra, y_vow, y_con]

Now the issue mostly I'm facing is to correctly define the metrics, loss, and loss_weights function for each of my outputs. However, I started as follows:

optimizer        = tf.keras.optimizers.Adam(learning_rate=0.05)
loss_fn          = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
train_acc_metric = tf.keras.metrics.Accuracy()

@tf.function
def train_step(x, y):
    with tf.GradientTape(persistent=True) as tape:
        logits = model(x, training=True)  # Logits for this minibatch
        train_loss_value = loss_fn(y, logits)

    grads = tape.gradient(train_loss_value, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))
    train_acc_metric.update_state(y, logits)
    return train_loss_value


for epoch in range(2):
    # Iterate over the batches of the dataset.
    for step, (x_batch_train, y_batch_train) in enumerate(train_generator):
        train_loss_value = train_step(x_batch_train, y_batch_train)

    # Reset metrics at the end of each epoch
    train_acc_metric.reset_states()

Apart from the above setup, I've tried other many ways to handle such problem cases though. For example, I defined 3 loss function and also 3 metrics as well but things are not working properly. The loss/acc became nan type stuff.

Here are my few straight queries in such case:

• how to define loss, metrics and loss_weights
• how to efficient use of all callbacks features

And just sake of learning opportunity, what if it has additionally regression type output (along with the rest 3 multi-out, so that total 4); how to deal all of them in custom fit? I've visited this SO, gave some hint for a different type of output (classification + regression).

Answer 1

You just need to do a custom training loop, but everything needs to be done 3 times (+ 1 if you also have a continuous variable). Here's an example using quadruple output architecture:

import tensorflow as tf
import numpy as np

(xtrain, train_target), (xtest, test_target) = tf.keras.datasets.mnist.load_data()

# 10 categories, one for each digit
ytrain1 = tf.keras.utils.to_categorical(train_target, num_classes=10)
ytest1 = tf.keras.utils.to_categorical(test_target, num_classes=10)

# 2 categories, if the digit is odd or not
ytrain2 = tf.keras.utils.to_categorical((train_target % 2 == 0).astype(int), 
                                        num_classes=2)
ytest2 = tf.keras.utils.to_categorical((test_target % 2 == 0).astype(int), 
                                       num_classes=2)

# 4 categories, based on the interval of the digit
ytrain3 = tf.keras.utils.to_categorical(np.digitize(train_target, [3, 6, 8]), 
                                        num_classes=4)
ytest3 = tf.keras.utils.to_categorical(np.digitize(test_target, [3, 6, 8]), 
                                       num_classes=4)

# Regression, the square of the digit
ytrain4 = tf.square(tf.cast(train_target, tf.float32))
ytest4 = tf.square(tf.cast(test_target, tf.float32))

# train dataset
train_ds = tf.data.Dataset. \
    from_tensor_slices((xtrain, ytrain1, ytrain2, ytrain3, ytrain4)). \
    shuffle(32). \
    batch(32).map(lambda a, *rest: (tf.divide(a[..., None], 255), rest)). \
    prefetch(tf.data.experimental.AUTOTUNE)

# test dataset
test_ds = tf.data.Dataset. \
    from_tensor_slices((xtest, ytest1, ytest2, ytest3, ytest4)). \
    shuffle(32). \
    batch(32).map(lambda a, *rest: (tf.divide(a[..., None], 255), rest)). \
    prefetch(tf.data.experimental.AUTOTUNE)


# architecture
class Net(tf.keras.Model):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = tf.keras.layers.Conv2D(filters=16, kernel_size=(3, 3),
                                            strides=(1, 1), input_shape=(28, 28, 1),
                                            activation='relu')
        self.maxp1 = tf.keras.layers.MaxPool2D(pool_size=(2, 2))
        self.conv2 = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3),
                                            strides=(1, 1),
                                            activation='relu')
        self.maxp2 = tf.keras.layers.MaxPool2D(pool_size=(2, 2))
        self.conv3 = tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3),
                                            strides=(1, 1),
                                            activation='relu')
        self.maxp3 = tf.keras.layers.MaxPool2D(pool_size=(2, 2))
        self.gap = tf.keras.layers.Flatten()
        self.dense = tf.keras.layers.Dense(64, activation='relu')
        self.output1 = tf.keras.layers.Dense(10, activation='softmax')
        self.output2 = tf.keras.layers.Dense(2, activation='softmax')
        self.output3 = tf.keras.layers.Dense(4, activation='softmax')
        self.output4 = tf.keras.layers.Dense(1, activation='linear')

    def call(self, inputs, training=False, **kwargs):
        x = self.conv1(inputs)
        x = self.maxp1(x)
        x = self.conv2(x)
        x = self.maxp2(x)
        x = self.conv3(x)
        x = self.maxp3(x)
        x = self.gap(x)
        x = self.dense(x)
        out1 = self.output1(x)
        out2 = self.output2(x)
        out3 = self.output3(x)
        out4 = self.output4(x)
        return out1, out2, out3, out4


model = Net()

optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

# the three losses
loss_1 = tf.losses.CategoricalCrossentropy()
loss_2 = tf.losses.CategoricalCrossentropy()
loss_3 = tf.losses.CategoricalCrossentropy()
loss_4 = tf.losses.MeanAbsoluteError()

# mean object that keeps track of the train losses
loss_1_train = tf.metrics.Mean(name='tr_loss_1')
loss_2_train = tf.metrics.Mean(name='tr_loss_2')
loss_3_train = tf.metrics.Mean(name='tr_loss_3')
loss_4_train = tf.metrics.Mean(name='tr_loss_4')

# mean object that keeps track of the test losses
loss_1_test = tf.metrics.Mean(name='ts_loss_1')
loss_2_test = tf.metrics.Mean(name='ts_loss_2')
loss_3_test = tf.metrics.Mean(name='ts_loss_3')
loss_4_test = tf.metrics.Mean(name='ts_loss_4')

# accuracies for printout
acc_1_train = tf.metrics.CategoricalAccuracy(name='tr_acc_1')
acc_2_train = tf.metrics.CategoricalAccuracy(name='tr_acc_2')
acc_3_train = tf.metrics.CategoricalAccuracy(name='tr_acc_3')

# accuracies for printout
acc_1_test = tf.metrics.CategoricalAccuracy(name='ts_acc_1')
acc_2_test = tf.metrics.CategoricalAccuracy(name='ts_acc_2')
acc_3_test = tf.metrics.CategoricalAccuracy(name='ts_acc_3')


# custom training loop
@tf.function
def train_step(x, y1, y2, y3, y4):
    with tf.GradientTape(persistent=True) as tape:
        out1, out2, out3, out4 = model(x, training=True)
        loss_1_value = loss_1(y1, out1)
        loss_2_value = loss_2(y2, out2)
        loss_3_value = loss_3(y3, out3)
        loss_4_value = loss_4(y4, out4)

    losses = [loss_1_value, loss_2_value, loss_3_value, loss_4_value]

    # a list of losses is passed
    grads = tape.gradient(losses, model.trainable_variables)

    # gradients are applied
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

    # losses are updated
    loss_1_train(loss_1_value)
    loss_2_train(loss_2_value)
    loss_3_train(loss_3_value)
    loss_4_train(loss_4_value)

    # accuracies are updated
    acc_1_train.update_state(y1, out1)
    acc_2_train.update_state(y2, out2)
    acc_3_train.update_state(y3, out3)


@tf.function
def test_step(x, y1, y2, y3, y4):
    out1, out2, out3, out4 = model(x, training=False)
    loss_1_value = loss_1(y1, out1)
    loss_2_value = loss_2(y2, out2)
    loss_3_value = loss_3(y3, out3)
    loss_4_value = loss_4(y4, out4)

    loss_1_test(loss_1_value)
    loss_2_test(loss_2_value)
    loss_3_test(loss_3_value)
    loss_4_test(loss_4_value)

    acc_1_test.update_state(y1, out1)
    acc_2_test.update_state(y2, out2)
    acc_3_test.update_state(y3, out3)


for epoch in range(5):
    # train step
    for inputs, outputs1, outputs2, outputs3, outputs4 in train_ds:
        train_step(inputs, outputs1, outputs2, outputs3, outputs4)

    # test step
    for inputs, outputs1, outputs2, outputs3, outputs4 in test_ds:
        test_step(inputs, outputs1, outputs2, outputs3, outputs4)

    metrics = [acc_1_train, acc_1_test,
               acc_2_train, acc_2_test,
               acc_3_train, acc_3_test,
               loss_4_train, loss_4_test]

    # printing metrics
    for metric in metrics:
        print(f'{metric.name}:{metric.result():=6.4f}', end=' ')   
    print()

    # resetting the states of the metrics
    loss_1_train.reset_states()
    loss_2_train.reset_states()
    loss_3_train.reset_states()

    loss_1_test.reset_states()
    loss_2_test.reset_states()
    loss_3_test.reset_states()

    acc_1_train.reset_states()
    acc_2_train.reset_states()
    acc_3_train.reset_states()

    acc_1_test.reset_states()
    acc_2_test.reset_states()
    acc_3_test.reset_states()

ts_acc_1:0.9495 ts_acc_2:0.9685 ts_acc_3:0.9589 ts_loss_4:5.5617 
ts_acc_1:0.9628 ts_acc_2:0.9747 ts_acc_3:0.9697 ts_loss_4:4.8953 
ts_acc_1:0.9697 ts_acc_2:0.9758 ts_acc_3:0.9733 ts_loss_4:4.5209 
ts_acc_1:0.9715 ts_acc_2:0.9796 ts_acc_3:0.9745 ts_loss_4:4.2175 
ts_acc_1:0.9742 ts_acc_2:0.9834 ts_acc_3:0.9775 ts_loss_4:3.9825

I wouldn't know how to use Keras Callbacks in a custom training loop, and neither does the most popular question on this topic. If you're looking to use EarlyStopping, I personally use a collections.deque, and interrupt when the minimum loss is the nth last. Here's an example:

from collections import deque
import numpy as np

epochs = 100
early_stopping = 5

loss_hist = deque(maxlen=early_stopping)

for epoch in range(epochs):
    loss_value = np.random.rand()
    loss_hist.append(loss_value)

    print('Last 5 values: ', *np.round(loss_hist, 3))

    if len(loss_hist) == early_stopping and loss_hist.popleft() < min(loss_hist):
        print('Early stopping. No loss decrease in %i epochs.\n' % early_stopping)
        break

Last 5 values:  0.456
Last 5 values:  0.456 0.153
Last 5 values:  0.456 0.153 0.2
Last 5 values:  0.456 0.153 0.2 0.433
Last 5 values:  0.456 0.153 0.2 0.433 0.528
Last 5 values:  0.153 0.2 0.433 0.528 0.349
Early stopping. No loss decrease in 5 epochs.

You can see that at the last time, the inner most value is the smallest of all, so there is no increase in validation loss. And that's the stopping condition.