Fix: move additional metrics from approximator to networks

[vpratz]

Supplying the additional metrics for inference and summary networks via the approximators compile method caused problems during deseralization (#497). This can be resolved nicely by moving the metrics directly to the networks' constructors, analogous to how Keras normally handles custom metrics in layers.

As summary networks and inference networks inherit from the respective base classes, this change only requires minor adaptations. Calls to layer_kwargs are now only used in classes that directly inherit from keras.Layer, and have been moved into InferenceNetwork and SummaryNetwork.

Fixes #497.

[codecov]

Codecov Report

Attention: Patch coverage is 88.05970% with 8 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...low/approximators/model_comparison_approximator.py	66.66%	3 Missing ⚠️
bayesflow/utils/serialization.py	92.30%	3 Missing ⚠️
bayesflow/approximators/continuous_approximator.py	66.66%	2 Missing ⚠️

Files with missing lines	Coverage Δ
...ow/experimental/diffusion_model/diffusion_model.py	79.03% <ø> (-0.12%)	⬇️
...w/networks/consistency_models/consistency_model.py	97.63% <100.00%> (-0.02%)	⬇️
bayesflow/networks/coupling_flow/coupling_flow.py	100.00% <ø> (ø)
bayesflow/networks/flow_matching/flow_matching.py	95.49% <ø> (-0.05%)	⬇️
bayesflow/networks/inference_network.py	78.26% <100.00%> (+2.07%)	⬆️
bayesflow/networks/mlp/mlp.py	90.90% <100.00%> (+0.16%)	⬆️
bayesflow/networks/summary_network.py	97.14% <100.00%> (+2.85%)	⬆️
bayesflow/approximators/continuous_approximator.py	85.00% <66.66%> (+0.17%)	⬆️
...low/approximators/model_comparison_approximator.py	81.39% <66.66%> (-0.43%)	⬇️
bayesflow/utils/serialization.py	90.58% <92.30%> (+1.22%)	⬆️

[vpratz]

@LarsKue @stefanradev93 @han-ol I encountered the problem that to pass metrics to the networks in their constructors, we would have to specify all get_config functions manually, as the metrics are not basic types.
I have now instead adapted the auto-config capabilities of keras.Layer to accept any object we can (de)serialize, and to use our serialization functions. This should make the auto-config functionality much more flexible. Take a look at the BaseLayer class for the implementation. It seems to me that this doesn't cause problems for now (hoping that all tests will pass).
Do you see any downsides with this approach? Or do you know something about the motivation of Keras to limit this functionality to basic types?

[stefanradev93]

I would let @LarsKue chip in on this, as I am bit concerned about having more and more "custom" variants of basic keras containers (e.g., Sequential, Layer,...).

[vpratz]

Thanks for the comment, I can understand this, relying too much on hacking/modifying Keras internals has the downside that our code might become more fragile with respect to changes in Keras. I think the ability to pass non-basic types to our inference and summary network base classes would be nice to have, the behavior in this PR is one example of this.
The options that I see are:

1. accept that we cannot pass non-basic arguments to the base classes, making implementation of some features more cumbersome (or impossible)
2. forego the auto-config from offered by keras.Layer and explicitly implement get_config in all our networks, which works but requires some redundant work and can be error-prone
3. a mechanism like the one implemented here, which offers more flexibility, but relies on Keras making not too drastic changes to the keras.Layer internals, which we cannot control

[LarsKue]

I would be on board if it was a lightweight wrapper, but it seems that there is a lot of copied code from keras, which we should avoid in general, imo.

@vpratz Can we solve this through monkey-patching somehow, like with the serialize and deserialize functions?

[stefanradev93]

If this is not breaking in any downstream task and only adds to the existing functionality, we may consider doing a PR to keras or asking them to implement it (or why they decided not to)?

[vpratz]

@LarsKue Thanks for taking a look. I'll take a look at alternative ways to achieve this behavior.

[vpratz]

@LarsKue Thanks a lot, this was a really good pointer. I have now implemented a wrapper around the constructor that is applied inside the serializable decorator. In addition, the default from_config is replaced with a variant that deserializes the config again.

There is still some copied code, but I think it is now limited to the part that is required for the feature to behave as similar to the Keras implementation as possible, apart from the desired changes (we might not need the part regarding the dtype, but I'm not sure and would leave it in there for now).

With those changes, we could remove the from_config method from most classes. Is this a change I should make?

@stefanradev93 @LarsKue Please take another look and let me know what you think (the most relevant part is in bayesflow/utils/serialization.py

[LarsKue]

Thanks, this looks much better. Since this is a sensitive change, I think we should extensively test it before rolling it out. Otherwise, green light from my side!

[vpratz]

I forgot to make the same changes in the ModelComparisonApproximator, this is now fixed. I have changed the output that is supplied to its classifier metrics from predictions to probabilities. This way, it is compatible with the metrics offered by Keras.