Query while building one Model in tf2

I am stucking with one problem while defining one model in tf2. I am using one dense layer with a softmax activation function. Now I want to extract the index with the highest probability value for that softmax layer output . So that I can use that index for later layer definition while building the model. Please help me out to implement the above problem, Waiting for your quick reply.

tf.math.argmax is the tensorflow operations which will allow you to do this. If you have a Keras model you can use a Custom layers, as described in the tutorial: Custom layers  |  TensorFlow Core

The final code will look something like:

class IndexOfMaxLayer(tf.keras.layers.Layer):
  def __init__(self):
    super(IndexOfMaxLayer, self).__init__()

  def call(self, inputs):
    return tf.math.argmax(inputs)

Now you can use this code after whichever dense / softmax layer you want to extract the max value from.

I want to select one layer from one list (5 layers are already stacked inside this list) for this corresponding index. But this above custom layer returns one keras tensor which does not help to select one layer from that list.

I didn’t really get your question, so you have a neural network, with 5 different outputs, and now you want to take the one with the presumably max probability, and propagate it’s outputs ahead? In this case it’s pretty much the same, except that you will first have to take the max of the outputs of each of those layers, put them together in a tensor, take the argmax of that tensor, and use it to select which layer you want to propagate from. I am not sure I get your problem fully, an example would help.

Thanks for your response… I am just giving example.

import tensorflow as tf

from transformers import BertTokenizerFast

from transformers import TFBertModel

mlm_model = TFBertModel.from_pretrained(‘bert-base-uncased’,output_attentions=True)

input_word_ids = tf.keras.layers.Input(shape=(50,), dtype=tf.int32,name=“input_word_ids”)

token_type_ids = tf.keras.layers.Input(shape=(50,), dtype=tf.int32,name=“token_type_ids”)

attention_mask = tf.keras.layers.Input(shape=(50,), dtype=tf.int32,name=“attention_mask”)

mlm_out=mlm_model([input_word_ids,token_type_ids,attention_mask])

print(“Attention Heads are”,mlm_out[2])

This mlm_out[2] has 12 attention heads now I want to extract the one attention head from mlm_out[2] which will have the highest probability. Is it possible or not …