Finally, we need the logits layer, which will take the output of the fully connected layer and then produce the raw prediction values. For example, in the case of the digit classification, the output will be a tensor of 10 values, where each value represents the score of one class from 0-9. So, let's define this logit layer for the digit classification example, where we need 10 outputs only, and with linear activation, which is the default for the dense() function of TensorFlow:

logits_layer = tf.layers.dense(inputs=dense_layer, units=10)

The final output of this logits layer will be a tensor of the following shape:

[batch_size, 10]

As mentioned previously, the logits layer of the model will return the raw predictions our our batch. But we need to convert these values to interpretable format:

• The predicted class for the input sample 0-9.
• The scores or probabilities for each possible class. For example, the probability that the sample is 0, is 1, and so on.

So, our predicted class will be the one that has the highest value in the 10 probabilities. We can get this value by using the argmax function as follows:

tf.argmax(input=logits_layer, axis=1)

Remember that the logits_layer has the following shape:

[batch_size, 10]

So, we need to find the max values along our predictions, which is the dimension that has an index of 1.

Finally, we can get our next value, which represents the probabilities of each target class, by applying softmax activation to the output of the logits_layer, which will squash each value to be between 0 and 1:

tf.nn.softmax(logits_layer, name="softmax_tensor")