|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| """A sample model implementation.
|
|
|
| This is only a dummy example to showcase how a model is composed. It is usually
|
| not needed to implement a model from scratch. Most SoTA models can be found and
|
| directly used from `official/vision/modeling` directory.
|
| """
|
|
|
| from typing import Any, Mapping
|
|
|
| import tensorflow as tf, tf_keras
|
| from official.vision.examples.starter import example_config as example_cfg
|
|
|
|
|
| class ExampleModel(tf_keras.Model):
|
| """A example model class.
|
|
|
| A model is a subclass of tf_keras.Model where layers are built in the
|
| constructor.
|
| """
|
|
|
| def __init__(
|
| self,
|
| num_classes: int,
|
| input_specs: tf_keras.layers.InputSpec = tf_keras.layers.InputSpec(
|
| shape=[None, None, None, 3]),
|
| **kwargs):
|
| """Initializes the example model.
|
|
|
| All layers are defined in the constructor, and config is recorded in the
|
| `_config_dict` object for serialization.
|
|
|
| Args:
|
| num_classes: The number of classes in classification task.
|
| input_specs: A `tf_keras.layers.InputSpec` spec of the input tensor.
|
| **kwargs: Additional keyword arguments to be passed.
|
| """
|
| inputs = tf_keras.Input(shape=input_specs.shape[1:], name=input_specs.name)
|
| outputs = tf_keras.layers.Conv2D(
|
| filters=16, kernel_size=3, strides=2, padding='same', use_bias=False)(
|
| inputs)
|
| outputs = tf_keras.layers.Conv2D(
|
| filters=32, kernel_size=3, strides=2, padding='same', use_bias=False)(
|
| outputs)
|
| outputs = tf_keras.layers.Conv2D(
|
| filters=64, kernel_size=3, strides=2, padding='same', use_bias=False)(
|
| outputs)
|
| outputs = tf_keras.layers.GlobalAveragePooling2D()(outputs)
|
| outputs = tf_keras.layers.Dense(1024, activation='relu')(outputs)
|
| outputs = tf_keras.layers.Dense(num_classes)(outputs)
|
|
|
| super().__init__(inputs=inputs, outputs=outputs, **kwargs)
|
| self._input_specs = input_specs
|
| self._config_dict = {'num_classes': num_classes, 'input_specs': input_specs}
|
|
|
| def get_config(self) -> Mapping[str, Any]:
|
| """Gets the config of this model."""
|
| return self._config_dict
|
|
|
| @classmethod
|
| def from_config(cls, config, custom_objects=None):
|
| """Constructs an instance of this model from input config."""
|
| return cls(**config)
|
|
|
|
|
| def build_example_model(input_specs: tf_keras.layers.InputSpec,
|
| model_config: example_cfg.ExampleModel,
|
| **kwargs) -> tf_keras.Model:
|
| """Builds and returns the example model.
|
|
|
| This function is the main entry point to build a model. Commonly, it builds a
|
| model by building a backbone, decoder and head. An example of building a
|
| classification model is at
|
| third_party/tensorflow_models/official/vision/modeling/backbones/resnet.py.
|
| However, it is not mandatory for all models to have these three pieces
|
| exactly. Depending on the task, model can be as simple as the example model
|
| here or more complex, such as multi-head architecture.
|
|
|
| Args:
|
| input_specs: The specs of the input layer that defines input size.
|
| model_config: The config containing parameters to build a model.
|
| **kwargs: Additional keyword arguments to be passed.
|
|
|
| Returns:
|
| A tf_keras.Model object.
|
| """
|
| return ExampleModel(
|
| num_classes=model_config.num_classes, input_specs=input_specs, **kwargs)
|
|
|
