I am trying to improve this situation, in image classification[3 classes, softmax in the last layer], I constructed the neural network having 7[conv2d+Batchnormalization] layers + 1 linear layer,
like[Conv2d, batch_normaliztion, Conv2d, batch_normaliztion, ..., Conv2d, batch_normaliztion, Dense layer]
And, this is the gradient flow for the last epoch
Here, I couldnt capture the Dense layer, where the gradient is changing in medium flow[average gradient for the Dense layer/Output layer = 0.03-0.07],
And, this is the loss graph
I am having trouble understanding that is this model - vanishing or exploding gradients, because the average gradients in all batch_normalization layers are changing but, the average gradients in the conv2d layers are not changing from first to last epoch.
Any ideas - how to improve this, and how to reduce the loss