Thanks to the contributions of @Pedro Marques and @lescurel, I managed to make a metric class that smooths any other metric across the epochs. It is also able to work when using a validation set during training thanks to the usage of different state variables.
There is a base class that manages the state logic:
class SmoothMetric(tf.keras.metrics.Metric):
def __init__(self, metric, name=None, **kwargs):
if name is None:
name = 'smooth_' + metric.name
super().__init__(name=name, **kwargs)
self.metric = metric
self.in_test_step = tf.Variable(False)
self.prev_in_test_step = tf.Variable(False)
self.train_states = {}
self.test_states = {}
def add_state_variable(self, name, **kwargs):
self.train_states[name] = tf.Variable(**kwargs)
self.test_states[name] = tf.Variable(**kwargs)
def get_state_variable(self, name):
return tf.cond(self.in_test_step, lambda: self.test_states[name], lambda: self.train_states[name])
def update_state(self, *args, **kwargs):
self.metric.update_state(*args, **kwargs)
self.prev_in_test_step.assign(self.in_test_step)
def update_filter_state_prevoius_epoch(self):
raise NotImplementedError('Needs to be overwritten')
def reset_state(self, *args, **kwargs):
# update state of the previous mode
tmp = self.in_test_step.read_value()
self.in_test_step.assign(self.prev_in_test_step)
self.update_filter_state_prevoius_epoch()
self.in_test_step.assign(tmp)
self.metric.reset_state(*args, **kwargs)
Then, I define a class that inherits from the previous, which implements the filter. In this case, it implements the Exponential Moving Average, but it is possible to make other filters such as the Simple Moving Average.
class SmoothMetricEMA(SmoothMetric):
def __init__(self, metric, alpha=0.1, name=None, **kwargs) -> None:
super().__init__(metric=metric, name=name, **kwargs)
self.alpha = alpha
self.add_state_variable('accum', initial_value=0.0)
self.add_state_variable('idx', initial_value=0)
def update_filter_state_prevoius_epoch(self):
self.get_state_variable('accum').assign(self.result())
self.get_state_variable('idx').assign_add(1)
def _result(self):
return self.metric.result()
def _filtered_result(self):
return (1 - self.alpha) * self.get_state_variable('accum') + self.alpha * self._result()
def result(self):
# if the idx is zero, the accumulator is 0 and _filtered_result will not give the correct result
return tf.cond(tf.equal(self.get_state_variable('idx'), 0), self._result, self._filtered_result)
Here is the SMA filter alternative:
class SmoothMetricSMA(SmoothMetric):
def __init__(self, metric, n_samples=10, remove_outliers=None, name=None, **kwargs) -> None:
super().__init__(metric=metric, name=name, **kwargs)
self.queue_size = n_samples - 1
self.remove_outliers = remove_outliers
self.add_state_variable('accum', initial_value=tf.zeros(self.queue_size))
self.add_state_variable('idx', initial_value=0)
def update_filter_state_prevoius_epoch(self):
self.get_state_variable('accum').scatter_update(
tf.IndexedSlices(self.metric.result(), self.get_state_variable('idx') % self.queue_size))
self.get_state_variable('idx').assign_add(1)
def _result(self):
return self.metric.result()
def _filtered_result(self):
queue = tf.concat([self.get_state_variable('accum')[:self.get_state_variable('idx')], tf.reshape(self._result(), (1,))], axis=0)
if self.remove_outliers is not None:
queue = tf.sort(queue)
eff_size = tf.cast(tf.minimum(self.get_state_variable('idx')+1, self.queue_size), tf.float32)
queue = queue[tf.cast(tf.math.floor(eff_size*self.remove_outliers/2), tf.int32):tf.cast(tf.math.ceil(eff_size*(1-self.remove_outliers/2)), tf.int32)]
return tf.reduce_mean(queue)
def result(self):
# if the idx is zero, the accumulator is 0 and _filtered_result will not give the correct result
return tf.cond(tf.equal(self.get_state_variable('idx'), 0), self._result, self._filtered_result)
Finally, for the filter to work, it is necessary to use the following callback. This tells the filter whether it is in training or test mode.
class NotifySmoothMetricsCallback(tf.keras.callbacks.Callback):
"""
Notifies the smooth metric if it is on train or test step.
"""
def on_test_begin(self, logs):
for metric in self.model.metrics:
if isinstance(metric, SmoothMetric):
metric.in_test_step.assign(True)
def on_test_end(self, logs):
for metric in self.model.metrics:
if isinstance(metric, SmoothMetric):
metric.in_test_step.assign(False)
def on_epoch_begin(self, epoch, logs):
for metric in self.model.metrics:
if isinstance(metric, SmoothMetric):
metric.in_test_step.assign(False)
def on_epoch_end(self, epoch, logs):
for metric in self.model.metrics:
if isinstance(metric, SmoothMetric):
metric.in_test_step.assign(False)
def on_train_begin(self, logs):
for metric in self.model.metrics:
if isinstance(metric, SmoothMetric):
metric.in_test_step.assign(False)
Here is an example of how to use these classes:
metrics=[
SmoothMetricEMA(tf.metrics.MeanSquaredError(), alpha=0.1, name='smooth'),
]
callbacks = [
tf.keras.callbacks.EarlyStopping(monitor='val_smooth', patience=10, verbose=1),
NotifySmoothMetricsCallback(),
]
model.compile(optimizer='adam', loss='mse', metrics=metrics)
model.fit(x_train, y_train, batch_size=32, epochs=100,
validation_data=(x_val, y_val),
callbacks=callbacks,
)
