Add a simple get_model for finetuning and easy of use.

Author: mingxingtan

efficientnetv2/effnetv2_model.py

@@ -459,13 +459,6 @@ def __init__(self, mconfig, name=None):
 
     self._avg_pooling = tf.keras.layers.GlobalAveragePooling2D(
         data_format=mconfig.data_format)
-    if mconfig.num_classes:
-      self._fc = tf.keras.layers.Dense(
-          mconfig.num_classes,
-          kernel_initializer=dense_kernel_initializer,
-          bias_initializer=tf.constant_initializer(mconfig.headbias or 0))
-    else:
-      self._fc = None
 
     if mconfig.dropout_rate > 0:
       self._dropout = tf.keras.layers.Dropout(mconfig.dropout_rate)
@@ -498,9 +491,6 @@ def call(self, inputs, training):
       self.endpoints['pooled_features'] = outputs
       if self._dropout:
         outputs = self._dropout(outputs, training=training)
-      self.endpoints['global_pool'] = outputs
-      if self._fc:
-        outputs = self._fc(outputs)
       self.endpoints['head'] = outputs
     return outputs
 
@@ -514,12 +504,13 @@ class EffNetV2Model(tf.keras.Model):
   def __init__(self,
                model_name='efficientnetv2-s',
                model_config=None,
+               include_top=True,
                name=None):
     """Initializes an `Model` instance.
 
     Args:
       model_name: A string of model name.
-      model_config: A dict of model configureations or a string of hparams.
+      model_config: A dict of model configurations or a string of hparams.
       name: A string of layer name.
 
     Raises:
@@ -533,6 +524,7 @@ def __init__(self,
     self.cfg = cfg
     self._mconfig = cfg.model
     self.endpoints = None
+    self.include_top = include_top
     self._build()
 
   def _build(self):
@@ -574,12 +566,25 @@ def _build(self):
     # Head part.
     self._head = Head(self._mconfig)
 
+    # top part for classification
+    if self.include_top and self._mconfig.num_classes:
+      self._fc = tf.keras.layers.Dense(
+          self._mconfig.num_classes,
+          kernel_initializer=dense_kernel_initializer,
+          bias_initializer=tf.constant_initializer(self._mconfig.headbias or 0))
+    else:
+      self._fc = None
+
   def summary(self, input_shape=(224, 224, 3), **kargs):
     x = tf.keras.Input(shape=input_shape)
     model = tf.keras.Model(inputs=[x], outputs=self.call(x, training=True))
     return model.summary()
 
-  def call(self, inputs, training, features_only=None, single_out=None):
+  def get_model_with_inputs(self, inputs, **kargs):
+    model = tf.keras.Model(inputs=[inputs], outputs=self.call(inputs, training=True))
+    return  model
+
+  def call(self, inputs, training, with_endpoints=False):
     """Implementation of call().
 
     Args:
@@ -624,19 +629,70 @@ def call(self, inputs, training, features_only=None, single_out=None):
             self.endpoints['reduction_%s/%s' % (reduction_idx, k)] = v
     self.endpoints['features'] = outputs
 
-    if not features_only:
-      # Calls final layers and returns logits.
-      outputs = self._head(outputs, training)
-      self.endpoints.update(self._head.endpoints)
-
-    if single_out:  # Use for building sequential models.
-      return outputs
-
-    return [outputs] + list(
-        filter(lambda endpoint: endpoint is not None, [
-            self.endpoints.get('reduction_1'),
-            self.endpoints.get('reduction_2'),
-            self.endpoints.get('reduction_3'),
-            self.endpoints.get('reduction_4'),
-            self.endpoints.get('reduction_5'),
-        ]))
+    # Head to obtain the final feature.
+    outputs = self._head(outputs, training)
+    self.endpoints.update(self._head.endpoints)
+
+    # Calls final dense layers and returns logits.
+    if self._fc:
+      with tf.name_scope('head'):  # legacy
+        outputs = self._fc(outputs)
+
+    if with_endpoints:  # Use for building sequential models.
+      return [outputs] + list(
+          filter(lambda endpoint: endpoint is not None, [
+              self.endpoints.get('reduction_1'),
+              self.endpoints.get('reduction_2'),
+              self.endpoints.get('reduction_3'),
+              self.endpoints.get('reduction_4'),
+              self.endpoints.get('reduction_5'),
+          ]))
+
+    return outputs
+
+
+def get_model(model_name,
+              model_config=None,
+              include_top=True,
+              pretrained=True,
+              training=True,
+              with_endpoints=False,
+              **kargs):
+  """Get a EfficientNet V1 or V2 model instance.
+
+  This is a simply utility for finetuning or inference.
+
+  Args:
+    model_name: a string such as 'efficientnetv2-s' or 'efficientnet-b0'.
+    model_config: A dict of model configurations or a string of hparams.
+    include_top: whether to include the final dense layer for classification.
+    pretrained: if true, download the checkpoint. If string, load the ckpt.
+    training: If true, all model variables are trainable.
+    with_endpoints: whether to return all intermedia endpoints.
+
+  Returns:
+    A single tensor if with_endpoints if False; otherwise, a list of tensor.
+  """
+  net = EffNetV2Model(model_name, model_config, include_top)
+  net(tf.keras.Input(shape=(None, None, 3)),
+      training=training,
+      with_endpoints=with_endpoints)
+  if pretrained is True:
+    # download checkpoint and set pretrained path. Supported models include:
+    #    efficientnetv2-s, efficientnetv2-m, efficientnetv2-l,
+    #    efficientnetv2-b0, efficientnetv2-b1, efficientnetv2-b2, efficientnetv2-b3, 
+    #    efficientnet-b0, efficientnet-b1, efficientnet-b2, efficientnet-b3,
+    #    efficientnet-b4, efficientnet-b5, efficientnet-b6, efficientnet-b7, efficientnet-l2
+    # More V2 ckpts: https://github.com/google/automl/tree/master/efficientnetv2
+    # More V1 ckpts: https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet
+    url = f'https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/{model_name}.tgz'
+    pretrained_ckpt = tf.keras.utils.get_file(model_name, url, untar=True)
+  else:
+    pretrained_ckpt = pretrained
+
+  if pretrained_ckpt:
+    if tf.io.gfile.isdir(pretrained_ckpt):
+      pretrained_ckpt = tf.train.latest_checkpoint(pretrained_ckpt)
+    net.load_weights(pretrained_ckpt)
+
+  return net

efficientnetv2/infer.py

@@ -44,94 +44,79 @@ def define_flags():
   flags.DEFINE_string('export_dir', None, 'Export or saved model directory')
   flags.DEFINE_string('trace_file', '/tmp/a.trace', 'If set, dump trace file.')
   flags.DEFINE_integer('batch_size', 16, 'Batch size.')
+  flags.DEFINE_bool('mixed_precision', False, 'If True, use mixed precision.')
 
 
 def get_config(model_name, dataset_cfg, hparam_str=''):
   """Create a keras model for EffNetV2."""
   config = copy.deepcopy(effnetv2_configs.get_model_config(model_name))
-  config.override(datasets.get_dataset_config(dataset_cfg))
-  config.override(hparam_str)
+  config.update(datasets.get_dataset_config(dataset_cfg))
+  config.override(hparam_str, allow_new_keys=True)
   config.model.num_classes = config.data.num_classes
   return config
 
 
 def build_tf2_model():
   """Build the tf2 model."""
   tf.config.run_functions_eagerly(FLAGS.debug)
-  config = get_config(FLAGS.model_name, FLAGS.dataset_cfg, FLAGS.hparam_str)
-  if config.runtime.mixed_precision:
+  if FLAGS.mixed_precision:
     # Use 'mixed_float16' if running on GPUs.
     policy = tf.keras.mixed_precision.Policy('mixed_float16')
     tf.keras.mixed_precision.set_global_policy(policy)
 
-  model = effnetv2_model.EffNetV2Model(FLAGS.model_name, config.model)
-  # Use call (not build) to match the namescope: tensorflow issues/29576
-  model(tf.ones([1, 224, 224, 3]), False)
-  if FLAGS.model_dir:
-    ckpt = FLAGS.model_dir
-    if tf.io.gfile.isdir(ckpt):
-      ckpt = tf.train.latest_checkpoint(FLAGS.model_dir)
-    model.load_weights(ckpt)
+  model = effnetv2_model.get_model(
+    FLAGS.model_name, FLAGS.hparam_str, include_top=True, pretrained=FLAGS.model_dir or True)
   model.summary()
-
-  class ExportModel(tf.Module):
-    """Export a saved model."""
-
-    def __init__(self, model):
-      super().__init__()
-      self.model = model
-
-    @tf.function
-    def f(self, images):
-      return self.model(images, training=False)[0]
-
-  return ExportModel(model)
-
+  return model
 
 def tf2_eval_dataset():
   """Run TF2 benchmark and inference."""
-  export_model = build_tf2_model()
-  isize = FLAGS.image_size or export_model.model.cfg.eval.isize
+  model = build_tf2_model()
+  isize = FLAGS.image_size or model.cfg.eval.isize
 
   def preprocess_fn(features):
     features['image'] = preprocessing.preprocess_image(
         features['image'], isize, is_training=False)
     return features
 
+  @tf.function
+  def f(x):
+    return model(x)
+
   top1_acc = tf.keras.metrics.Accuracy()
   pbar = tf.keras.utils.Progbar(None)
   data = tfds.load('imagenet2012', split='validation')
   ds = data.map(preprocess_fn).batch(FLAGS.batch_size)
   for i, batch in enumerate(ds.prefetch(tf.data.experimental.AUTOTUNE)):
-    logits = export_model.f(batch['image'])
+    logits = f(batch['image'])
     top1_acc.update_state(batch['label'], tf.argmax(logits, axis=-1))
     pbar.update(i, [('top1', top1_acc.result().numpy())])
   print('\n top1= {:.4f}'.format(top1_acc.result().numpy()))
 
 
 def tf2_benchmark():
   """Run TF2 benchmark and inference."""
-  export_model = build_tf2_model()
-  isize = FLAGS.image_size or export_model.model.cfg.eval.isize
+  model = build_tf2_model()
+  isize = FLAGS.image_size or model.cfg.eval.isize
   if FLAGS.export_dir:
-    tf.saved_model.save(
-        export_model,
-        FLAGS.export_dir,
-        signatures=export_model.f.get_concrete_function(
-            tf.TensorSpec(shape=(None, isize, isize, 3), dtype=tf.float16)))
-    export_model = tf.saved_model.load(FLAGS.export_dir)
+    tf.saved_model.save(model, FLAGS.export_dir)
+    model = tf.saved_model.load(FLAGS.export_dir)
 
   batch_size = FLAGS.batch_size
   imgs = tf.ones((batch_size, isize, isize, 3), dtype=tf.float16)
 
+  @tf.function
+  def f(x):
+    return model(x)
+
   print('starting warmup.')
   for _ in range(10):  # warmup runs.
-    export_model.f(imgs)
+    f(imgs)
 
   print('start benchmark.')
   start = time.perf_counter()
   for _ in range(10):
-    export_model.f(imgs)
+    f(imgs)
   end = time.perf_counter()
   inference_time = (end - start) / 10
 
@@ -143,14 +128,13 @@ def tf1_benchmark():
   """Run TF1 inference and benchmark."""
   # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
   from tensorflow.python.client import timeline
-  config = get_config(FLAGS.model_name, FLAGS.dataset_cfg, FLAGS.hparam_str)
   with tf1.Session() as sess:
-    model = effnetv2_model.EffNetV2Model(FLAGS.model_name, config.model)
+    model = effnetv2_model.EffNetV2Model(FLAGS.model_name, FLAGS.hparam_str)
     batch_size = FLAGS.batch_size
     run_options = tf1.RunOptions(
         trace_level=tf1.RunOptions.FULL_TRACE)
     run_metadata = tf1.RunMetadata()
-    isize = FLAGS.image_size or config.eval.isize
+    isize = FLAGS.image_size or model.cfg.eval.isize
     inputs = tf.ones((batch_size, isize, isize, 3), tf.float16)
     output = model(inputs, training=False)
     sess.run(tf1.global_variables_initializer())
@@ -179,10 +163,9 @@ def tf1_benchmark():
 def tf1_export_ema_ckpt():
   """Restore variables from a given checkpoint."""
   with tf1.Session() as sess:
-    config = get_config(FLAGS.model_name, FLAGS.dataset_cfg, FLAGS.hparam_str)
-    model = effnetv2_model.EffNetV2Model(FLAGS.model_name, config.model)
+    model = effnetv2_model.EffNetV2Model(FLAGS.model_name, FLAGS.hparam_str)
     batch_size = FLAGS.batch_size
-    isize = FLAGS.image_size or config.eval.isize
+    isize = FLAGS.image_size or model.cfg.eval.isize
     inputs = tf.ones((batch_size, isize, isize, 3), tf.float32)
     _ = model(inputs, training=False)
     sess.run(tf1.global_variables_initializer())