Spaces:
Paused
Paused
| # Copyright 2022 Google LLC | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # https://www.apache.org/licenses/LICENSE-2.0 | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # ============================================================================== | |
| """A library for instantiating frame interpolation evaluation metrics.""" | |
| from typing import Callable, Dict, Text | |
| from ..losses import losses | |
| import tensorflow as tf | |
| class TrainLossMetric(tf.keras.metrics.Metric): | |
| """Compute training loss for our example and prediction format. | |
| The purpose of this is to ensure that we always include a loss that is exactly | |
| like the training loss into the evaluation in order to detect possible | |
| overfitting. | |
| """ | |
| def __init__(self, name='eval_loss', **kwargs): | |
| super(TrainLossMetric, self).__init__(name=name, **kwargs) | |
| self.acc = self.add_weight(name='train_metric_acc', initializer='zeros') | |
| self.count = self.add_weight(name='train_metric_count', initializer='zeros') | |
| def update_state(self, | |
| batch, | |
| predictions, | |
| sample_weight=None, | |
| checkpoint_step=0): | |
| loss_functions = losses.training_losses() | |
| loss_list = [] | |
| for (loss_value, loss_weight) in loss_functions.values(): | |
| loss_list.append( | |
| loss_value(batch, predictions) * loss_weight(checkpoint_step)) | |
| loss = tf.add_n(loss_list) | |
| self.acc.assign_add(loss) | |
| self.count.assign_add(1) | |
| def result(self): | |
| return self.acc / self.count | |
| def reset_states(self): | |
| self.acc.assign(0) | |
| self.count.assign(0) | |
| class L1Metric(tf.keras.metrics.Metric): | |
| """Compute L1 over our training example and prediction format. | |
| The purpose of this is to ensure that we have at least one metric that is | |
| compatible across all eval the session and allows us to quickly compare models | |
| against each other. | |
| """ | |
| def __init__(self, name='eval_loss', **kwargs): | |
| super(L1Metric, self).__init__(name=name, **kwargs) | |
| self.acc = self.add_weight(name='l1_metric_acc', initializer='zeros') | |
| self.count = self.add_weight(name='l1_metric_count', initializer='zeros') | |
| def update_state(self, batch, prediction, sample_weight=None, | |
| checkpoint_step=0): | |
| self.acc.assign_add(losses.l1_loss(batch, prediction)) | |
| self.count.assign_add(1) | |
| def result(self): | |
| return self.acc / self.count | |
| def reset_states(self): | |
| self.acc.assign(0) | |
| self.count.assign(0) | |
| class GenericLossMetric(tf.keras.metrics.Metric): | |
| """Metric based on any loss function.""" | |
| def __init__(self, name: str, loss: Callable[..., tf.Tensor], | |
| weight: Callable[..., tf.Tensor], **kwargs): | |
| """Initializes a metric based on a loss function and a weight schedule. | |
| Args: | |
| name: The name of the metric. | |
| loss: The callable loss that calculates a loss value for a (prediction, | |
| target) pair. | |
| weight: The callable weight scheduling function that samples a weight | |
| based on iteration. | |
| **kwargs: Any additional keyword arguments to be passed. | |
| """ | |
| super(GenericLossMetric, self).__init__(name=name, **kwargs) | |
| self.acc = self.add_weight(name='loss_metric_acc', initializer='zeros') | |
| self.count = self.add_weight(name='loss_metric_count', initializer='zeros') | |
| self.loss = loss | |
| self.weight = weight | |
| def update_state(self, | |
| batch, | |
| predictions, | |
| sample_weight=None, | |
| checkpoint_step=0): | |
| self.acc.assign_add( | |
| self.loss(batch, predictions) * self.weight(checkpoint_step)) | |
| self.count.assign_add(1) | |
| def result(self): | |
| return self.acc / self.count | |
| def reset_states(self): | |
| self.acc.assign(0) | |
| self.count.assign(0) | |
| def create_metrics_fn() -> Dict[Text, tf.keras.metrics.Metric]: | |
| """Create evaluation metrics. | |
| L1 and total training loss are added by default. | |
| The rest are the configured by the test_losses item via gin. | |
| Returns: | |
| A dictionary from metric name to Keras Metric object. | |
| """ | |
| metrics = {} | |
| # L1 is explicitly added just so we always have some consistent numbers around | |
| # to compare across sessions. | |
| metrics['l1'] = L1Metric() | |
| # We also always include training loss for the eval set to detect overfitting: | |
| metrics['training_loss'] = TrainLossMetric() | |
| test_losses = losses.test_losses() | |
| for loss_name, (loss_value, loss_weight) in test_losses.items(): | |
| metrics[loss_name] = GenericLossMetric( | |
| name=loss_name, loss=loss_value, weight=loss_weight) | |
| return metrics | |