Add warning message when deterministic training loss stagnates too quickly in partial BNNs

[sarah-allec]

I've encountered one dataset for which mini-batching in the deterministic neural network (DNN) training of a partial Bayesian neural network (PBNN) causes the DNN training loss to stagnate too quickly, resulting in significantly lower performance than non-Bayesian machine learning models. Increasing the batch size (or not doing any batching) alleviated the issue. It would be helpful if there was a check during DNN training that triggers a warning when the DNN training loss stagnates prematurely, offering a suggestion to increase the batch size as a possible solution.

[ziatdinovmax]

Thanks for bringing this up! We could potentially implement this as a utility function in neurobayes/utils that monitors the DNN training loss over a set number of epochs and triggers a warning if it detects that the loss has plateaued prematurely. This could include a suggestion to try changing the batch size, adjusting the MAP sigma value, or modifying the learning rate (the first two are the most frequent culprits in my experience).

Would you be open to submitting a PR with such a utility function?

[sarah-allec]

Yes, will do!