search

neuronets / nobrainer

A framework for developing neural network models for 3D image processing.
Other
158 stars 46 forks source link

Rewrite metrics module with model subclassing #343

Open hvgazula opened 1 month ago

hvgazula commented 1 month ago

What would you like changed/added and why? Rewrite functions in metrics.py with model subclassing. See here. The drawback of the current approach is that metrics have to be kept track of when training models and that can not always be possible.

What would be the benefit? Currently, loading a model from a previously stored checkpoint (warm_start) throws the following warning

WARNING:tensorflow:Unable to restore custom metric. Please ensure that the layer implements get_config and from_config when saving. In addition, please use the custom_objects arg when calling load_model().

Does the change make something easier to use? This change will avoid situations where the user has to document the metric used and subsequently pass it again at load time. Model sub-classing (with @keras.saving.register_keras_serializable) will greatly simplify this process.

hvgazula commented 1 month ago

NOTE: dice and tversky loss have been introduced in tf 2.16..