title: Multioutput-Multiclass Classification in Custom Scratch Training in TF.Keras
body: <p>I want to train a <strong>multi-out and multi-class</strong> classification model from scratch (using custom <code>fit()</code>). 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.</p>
  <h2>Data Set and Goal</h2>
  <p>I'm using data from <a href="https://www.kaggle.com/c/bengaliai-cv19" rel="nofollow noreferrer">here</a>; It's a Bengali handwritten character recognition challenge, each of the samples has <strong>3</strong> mutually related output along with <strong>multiple classes</strong> of each. Please see the figure below:</p>
  <p><img src="https://user-images.githubusercontent.com/17668390/94802661-e3eb5680-0409-11eb-8b43-76dd355acb89.png" alt="a" /></p>
  <p>In the above figure, as you can see, the ক্ট্রো is consist of 3 component (ক্ট , ো , ‍‍্র), namely <strong>Grapheme Root</strong>, <strong>Vowel Diactrics</strong> and <strong>Consonant Diacritics</strong> respectively and together they're called <strong>Grapheme</strong>. Again the <strong>Grapheme Root</strong> also has <strong>168</strong> different categories and also same as others (<strong>11</strong> and <strong>7</strong>). The added complexity results in <strong>~13,000</strong> different grapheme variations (compared to English’s 250 graphemic units).</p>
  <p>The goal is to classify the <strong>Components of the Grapheme</strong> in each image.</p>
  <h2>Initial Approach (and no issue with it)</h2>
  <p>I implemented a training pipeline over <a href="https://www.kaggle.com/ipythonx/keras-grapheme-gridmask-augmix-ensemble" rel="nofollow noreferrer">here</a>, where it's demonstrated using old <code>keras</code> (not <code>tf.keras</code>) with its a convenient feature such as <code>model.compile</code>, <code>callbacks</code> etc. I defined a <a href="https://www.kaggle.com/ipythonx/keras-grapheme-gridmask-augmix-ensemble#Grapheme-Data-Generator" rel="nofollow noreferrer">custom data generator</a> and defined a model architecture something like below.</p>
  <pre class="lang-py prettyprint-override"><code>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)
  </code></pre>
  <p>And compile the model as follows:</p>
  <pre class="lang-py prettyprint-override"><code>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'}
      )
  </code></pre>
  <p>As you can see I can cleary control each of the outputs with specific <code>loss</code>, <code>loss_weights</code>, and <code>accuracy</code>. And using the <code>.fit()</code> method, it's feasible to use any <code>callbacks</code> function for the model.</p>
  <h2>New Approach (and some issue with it)</h2>
  <p>Now, I want to re-implement it with the new feature of <code>tf.keras</code>. Such as <strong>model subclassing</strong> and <strong>custom fit training</strong>.  However, no change in the data loader. The model is defined as follows:</p>
  <pre class="lang-py prettyprint-override"><code>    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]
  </code></pre>
  <p>Now the issue mostly I'm facing is to correctly define the <code>metrics</code>, <code>loss</code>, and <code>loss_weights</code> function for each of my outputs. However, I started as follows:</p>
  <pre class="lang-py prettyprint-override"><code>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()
  </code></pre>
  <p>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 <code>loss/acc</code> became <code>nan</code> type stuff.</p>
  <p>Here are my few straight queries in such case:</p>
  <ul>
  <li>how to define <code>loss</code>, <code>metrics</code> and <code>loss_weights</code></li>
  <li>how to efficient use of all <code>callbacks</code> features</li>
  </ul>
  <p>And just sake of learning opportunity, what if it has additionally <strong>regression</strong> type output (along with the rest <strong>3</strong> multi-out, so that total <strong>4</strong>); how to deal all of them in custom <code>fit</code>? I've visited this <a href="https://stackoverflow.com/questions/59690188/how-do-i-make-a-multi-output-tensorflow-2-0-neural-network-with-two-different-va">SO</a>, gave some hint for a different type of output (<code>classification</code> + <code>regression</code>).</p>