How to control reduction strategy for stateful metric in keras mirrored strategy

user_4060812 0
Viewed 294

I use keras fit() method with custom metrics passed to model. The metrics are stateful - i.e. are a subclass of a Metric, as described in https://keras.io/api/metrics/#as-subclasses-of-metric-stateful When I run the code in a multi-gpu environment using a tf.distribute.MirroredStrategy() my metric code is called on every GPU separately with batch_size/no_of_gpus examples passed, which is reasonable to expect.

What happens next is that multiple scalars (one from every GPU) of the metric value need to be reduced to a single scalar, and what I get all the time is a sum reduction, while I would like to control that. Keep in mind, that reduction parameter is the one of Loss in keras, and there is no such thing in the Metric class: https://github.com/tensorflow/tensorflow/blob/acbc065f8eb2ed05c7ab5c42b5c5bd6abdd2f91f/tensorflow/python/keras/metrics.py#L87

(the only crazy thing I tried was to inherit from a Mean class that is a subclass of a Metric but that didn't change anything)

reduction is mentioned in the metrics code, however this is a reduction over multiple accumulated values in a single metric object, and in multi-gpu setting - this is not the case, as every metric works in its own GPU and is somehow aggregated at the end.

The way I debugged it to understand this behaviour was - I was printing the shapes and the results inside update_state method of the metric. And then I looked at value of the metric in logs object in on_batch_end callback.

I tried looking at TF code, but couldn't find the place this is happening. I would like to be able to control this behaviour - so either pick 'mean' or 'sum' for the metric, or at least know where it is being done in the code.

Edited: I guess this https://github.com/tensorflow/tensorflow/issues/39268 sheds some more light on this issue

multi-gpu keras tensorflow
1 Answers

I am facing the same problem as you (and that's why I found your question).

Seeing that it's been 15 days since you asked the question and there are no answers/comments yet, I thought I might share my temporary workaround.

Like you, I also think that a SUM reduction has been performed when combining progress over multiple GPUs. What I did is to pass the number of GPUs (e.g. given by the num_replicas_in_sync attribute of your tf.distribute strategy object) into the __init__(...) constructor of your sub-classed metric object, and use it to divide the return value in the results() method.

Potentially, you could also use tf.distribute.get_strategy() from within the metric object to make it "strategy aware", and use the information to decide how to modify the values in an ad hoc manner so that the SUM reduction will produce what you want.

I hope this helps for now, whether as a suggestion or as a confirmation that you're not alone on this.

user_129603680
Related