Spaces:
Build error
Build error
| # https://github.com/tensorflow/models/blob/master/official/vision/beta/modeling/classification_model.py | |
| """Build classification models.""" | |
| from typing import Any, Mapping, Optional, Union | |
| # Import libraries | |
| import tensorflow as tf | |
| import re | |
| class ClassificationModel(tf.keras.Model): | |
| """A classification class builder.""" | |
| def __init__( | |
| self, | |
| backbone: tf.keras.Model, | |
| num_classes: int, | |
| input_specs: tf.keras.layers.InputSpec = tf.keras.layers.InputSpec(shape=[None, None, None, 3]), | |
| dropout_rate: float = 0.0, | |
| kernel_initializer: Union[str, tf.keras.initializers.Initializer] = 'random_uniform', | |
| skip_logits_layer: bool = False, | |
| weight_decay: float = 0.0, | |
| clip_grad_norm: float = 0.0, | |
| **kwargs): | |
| """Classification initialization function. | |
| Args: | |
| backbone: a backbone network. | |
| num_classes: `int` number of classes in classification task. | |
| input_specs: `tf.keras.layers.InputSpec` specs of the input tensor. | |
| dropout_rate: `float` rate for dropout regularization. | |
| kernel_initializer: kernel initializer for the dense layer. | |
| kernel_regularizer: tf.keras.regularizers.Regularizer object. Default to | |
| None. | |
| bias_regularizer: tf.keras.regularizers.Regularizer object. Default to | |
| None. | |
| skip_logits_layer: `bool`, whether to skip the prediction layer. | |
| **kwargs: keyword arguments to be passed. | |
| """ | |
| inputs = tf.keras.Input(shape=input_specs.shape[1:], name=input_specs.name) | |
| outputs = backbone(inputs) | |
| outputs = tf.keras.layers.GlobalAveragePooling2D()(outputs) | |
| if not skip_logits_layer: | |
| if dropout_rate is not None and dropout_rate > 0: | |
| outputs = tf.keras.layers.Dropout(dropout_rate)(outputs) | |
| outputs = tf.keras.layers.Dense( | |
| num_classes, | |
| kernel_initializer=kernel_initializer, | |
| bias_initializer=tf.keras.initializers.Constant(value=-10.0))(outputs) | |
| super(ClassificationModel, self).__init__(inputs=inputs, outputs=outputs, **kwargs) | |
| self._config_dict = { | |
| 'backbone': backbone, | |
| 'num_classes': num_classes, | |
| 'input_specs': input_specs, | |
| 'dropout_rate': dropout_rate, | |
| 'kernel_initializer': kernel_initializer, | |
| 'weight_decay': weight_decay, | |
| 'clip_grad_norm': clip_grad_norm, | |
| } | |
| self._input_specs = input_specs | |
| self._backbone = backbone | |
| self._weight_decay = weight_decay | |
| self._clip_grad_norm = clip_grad_norm | |
| def _reg_l2_loss(self, weight_decay, regex=r'.*(kernel|weight):0$'): | |
| """Return regularization l2 loss loss.""" | |
| var_match = re.compile(regex) | |
| return weight_decay * tf.add_n([ | |
| tf.nn.l2_loss(v) | |
| for v in self.trainable_variables | |
| if var_match.match(v.name) | |
| ]) | |
| def train_step(self, data): | |
| features, labels = data | |
| images, labels = features['image'], labels['label'] | |
| with tf.GradientTape() as tape: | |
| pred = self(images, training=True) | |
| pred = tf.cast(pred, tf.float32) | |
| loss = self.compiled_loss( | |
| labels, | |
| pred, | |
| regularization_losses=[self._reg_l2_loss(self._weight_decay)]) | |
| self.optimizer.minimize(loss, self.trainable_variables, tape=tape) | |
| self.compiled_metrics.update_state(labels, pred) | |
| return {m.name: m.result() for m in self.metrics} | |
| def test_step(self, data): | |
| features, labels = data | |
| images, labels = features['image'], labels['label'] | |
| pred = self(images, training=False) | |
| pred = tf.cast(pred, tf.float32) | |
| self.compiled_loss( | |
| labels, | |
| pred, | |
| regularization_losses=[self._reg_l2_loss(self._weight_decay)]) | |
| self.compiled_metrics.update_state(labels, pred) | |
| return {m.name: m.result() for m in self.metrics} | |
| def checkpoint_items(self) -> Mapping[str, tf.keras.Model]: | |
| """Returns a dictionary of items to be additionally checkpointed.""" | |
| return dict(backbone=self.backbone) | |
| def backbone(self) -> tf.keras.Model: | |
| return self._backbone | |
| def get_config(self) -> Mapping[str, Any]: | |
| return self._config_dict | |
| def from_config(cls, config, custom_objects=None): | |
| return cls(**config) |