# Copyright 2023 The TensorFlow Authors. All Rights Reserved. # # 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 # # http://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. """An example task definition for image classification.""" from typing import Any, List, Optional, Tuple, Sequence, Mapping import tensorflow as tf, tf_keras from official.common import dataset_fn from official.core import base_task from official.core import task_factory from official.modeling import tf_utils from official.vision.dataloaders import input_reader_factory from official.vision.examples.starter import example_config as exp_cfg from official.vision.examples.starter import example_input from official.vision.examples.starter import example_model @task_factory.register_task_cls(exp_cfg.ExampleTask) class ExampleTask(base_task.Task): """Class of an example task. A task is a subclass of base_task.Task that defines model, input, loss, metric and one training and evaluation step, etc. """ def build_model(self) -> tf_keras.Model: """Builds a model.""" input_specs = tf_keras.layers.InputSpec(shape=[None] + self.task_config.model.input_size) model = example_model.build_example_model( input_specs=input_specs, model_config=self.task_config.model) return model def build_inputs( self, params: exp_cfg.ExampleDataConfig, input_context: Optional[tf.distribute.InputContext] = None ) -> tf.data.Dataset: """Builds input. The input from this function is a tf.data.Dataset that has gone through pre-processing steps, such as augmentation, batching, shuffling, etc. Args: params: The experiment config. input_context: An optional InputContext used by input reader. Returns: A tf.data.Dataset object. """ num_classes = self.task_config.model.num_classes input_size = self.task_config.model.input_size decoder = example_input.Decoder() parser = example_input.Parser( output_size=input_size[:2], num_classes=num_classes) reader = input_reader_factory.input_reader_generator( params, dataset_fn=dataset_fn.pick_dataset_fn(params.file_type), decoder_fn=decoder.decode, parser_fn=parser.parse_fn(params.is_training)) dataset = reader.read(input_context=input_context) return dataset def build_losses(self, labels: tf.Tensor, model_outputs: tf.Tensor, aux_losses: Optional[Any] = None) -> tf.Tensor: """Builds losses for training and validation. Args: labels: Input groundt-ruth labels. model_outputs: Output of the model. aux_losses: The auxiliarly loss tensors, i.e. `losses` in tf_keras.Model. Returns: The total loss tensor. """ total_loss = tf_keras.losses.sparse_categorical_crossentropy( labels, model_outputs, from_logits=True) total_loss = tf_utils.safe_mean(total_loss) if aux_losses: total_loss += tf.add_n(aux_losses) return total_loss def build_metrics(self, training: bool = True) -> Sequence[tf_keras.metrics.Metric]: """Gets streaming metrics for training/validation. This function builds and returns a list of metrics to compute during training and validation. The list contains objects of subclasses of tf_keras.metrics.Metric. Training and validation can have different metrics. Args: training: Whether the metric is for training or not. Returns: A list of tf_keras.metrics.Metric objects. """ k = self.task_config.evaluation.top_k metrics = [ tf_keras.metrics.SparseCategoricalAccuracy(name='accuracy'), tf_keras.metrics.SparseTopKCategoricalAccuracy( k=k, name='top_{}_accuracy'.format(k)) ] return metrics def train_step(self, inputs: Tuple[Any, Any], model: tf_keras.Model, optimizer: tf_keras.optimizers.Optimizer, metrics: Optional[List[Any]] = None) -> Mapping[str, Any]: """Does forward and backward. This example assumes input is a tuple of (features, labels), which follows the output from data loader, i.e., Parser. The output from Parser is fed into train_step to perform one step forward and backward pass. Other data structure, such as dictionary, can also be used, as long as it is consistent between output from Parser and input used here. Args: inputs: A tuple of input tensors of (features, labels). model: A tf_keras.Model instance. optimizer: The optimizer for this training step. metrics: A nested structure of metrics objects. Returns: A dictionary of logs. """ features, labels = inputs num_replicas = tf.distribute.get_strategy().num_replicas_in_sync with tf.GradientTape() as tape: outputs = model(features, training=True) # Casting output layer as float32 is necessary when mixed_precision is # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32. outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs) # Computes per-replica loss. loss = self.build_losses( model_outputs=outputs, labels=labels, aux_losses=model.losses) # Scales loss as the default gradients allreduce performs sum inside the # optimizer. scaled_loss = loss / num_replicas # For mixed_precision policy, when LossScaleOptimizer is used, loss is # scaled for numerical stability. if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer): scaled_loss = optimizer.get_scaled_loss(scaled_loss) tvars = model.trainable_variables grads = tape.gradient(scaled_loss, tvars) # Scales back gradient before apply_gradients when LossScaleOptimizer is # used. if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer): grads = optimizer.get_unscaled_gradients(grads) optimizer.apply_gradients(list(zip(grads, tvars))) logs = {self.loss: loss} if metrics: self.process_metrics(metrics, labels, outputs) return logs def validation_step(self, inputs: Tuple[Any, Any], model: tf_keras.Model, metrics: Optional[List[Any]] = None) -> Mapping[str, Any]: """Runs validation step. Args: inputs: A tuple of input tensors of (features, labels). model: A tf_keras.Model instance. metrics: A nested structure of metrics objects. Returns: A dictionary of logs. """ features, labels = inputs outputs = self.inference_step(features, model) outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs) loss = self.build_losses( model_outputs=outputs, labels=labels, aux_losses=model.losses) logs = {self.loss: loss} if metrics: self.process_metrics(metrics, labels, outputs) return logs def inference_step(self, inputs: tf.Tensor, model: tf_keras.Model) -> Any: """Performs the forward step. It is used in 'validation_step'.""" return model(inputs, training=False)