fix colab tpu training (#1050)

Author: fsx950223

efficientdet/utils.py

@@ -376,7 +376,7 @@ def __init__(self, name, _):  # pylint: disable=super-init-not-called
 
 def scalar(name, tensor, is_tpu=True):
   """Stores a (name, Tensor) tuple in a custom collection."""
-  logging.info('Adding scale summary {}'.format(Pair(name, tensor)))
+  logging.info('Adding scalar summary {}'.format(Pair(name, tensor)))
   if is_tpu:
     tf.add_to_collection('scalar_summaries', Pair(name, tf.reduce_mean(tensor)))
   else:

efficientnetv2/datasets.py

@@ -531,7 +531,7 @@ def _input_fn(self, batch_size, current_host, num_hosts):
     logging.info('use tfds: %s[%s]', self.cfg.tfds_name,
                  self.cfg.splits[self.split]['tfds_split'])
     ds = tfds.load(
-        self.cfg.tfds_name, split=self.cfg.splits[self.split]['tfds_split'])
+        self.cfg.tfds_name, split=self.cfg.splits[self.split]['tfds_split'], try_gcs=self.cfg.try_gcs)
     ds = ds.shard(num_hosts, current_host)
     if self.is_training:
       if self.cache:
@@ -581,6 +581,21 @@ class FlowersInput(CIFAR10Input):
           )))
 
 
+class TFFlowersInput(CIFAR10Input):
+  """TFFlowers input from tfds gcs."""
+  cfg = copy.deepcopy(CIFAR10Input.cfg)
+  cfg.update(
+      dict(
+          num_classes=5,
+          tfds_name='tf_flowers',
+          try_gcs=True,
+          splits=dict(
+              train=dict(num_images=2569, tfds_split='train[:70%]'),
+              minival=dict(num_images=1101, tfds_split='train[30%:]'),
+              eval=dict(num_images=1101, tfds_split='train[30%:]'),
+          )))
+
+
 class CarsInput(CIFAR10Input):
   """Car input from tfds."""
   cfg = copy.deepcopy(CIFAR10Input.cfg)
@@ -620,6 +635,7 @@ def get_dataset_class(ds_name):
       'cifar10': CIFAR10Input,
       'cifar100': CIFAR100Input,
       'flowers': FlowersInput,
+      'tfflowers': TFFlowersInput,
       'cars': CarsInput,
   }[ds_name]
 
@@ -730,6 +746,11 @@ class FlowersFt(Cifar10Ft):
   cfg = copy.deepcopy(Cifar10Ft.cfg)
   cfg.data.override(dict(ds_name='flowers'))
 
+@ds_register
+class TFFlowersFt(Cifar10Ft):
+  """Finetune tfflower configs."""
+  cfg = copy.deepcopy(Cifar10Ft.cfg)
+  cfg.data.override(dict(ds_name='tfflowers'))
 
 @ds_register
 class CarsFt(Cifar10Ft):

efficientnetv2/hparams.py

@@ -288,6 +288,7 @@ def lookup(name, prefix='effnet:') -> Any:
         multiclass=None,
         num_classes=1000,
         tfds_name=None,
+        try_gcs=False,
         tfds_split=None,
         splits=dict(
             train=dict(

efficientnetv2/main_tf2.py

@@ -214,6 +214,11 @@ def main(_) -> None:
         log_dir=FLAGS.model_dir, update_freq=100)
     rstr_callback = utils.ReuableBackupAndRestore(backup_dir=FLAGS.model_dir)
 
+    def filter_callbacks(callbacks):
+      if strategy == 'tpu' and not FLAGS.model_dir.startswith('gs://'):
+        return list(filter(lambda callback: isinstance(callback, tf.keras.callbacks.ModelCheckpoint), callbacks))
+      return callbacks
+
     def get_dataset(training, image_size, config):
       """A shared utility to get input dataset."""
       if training:
@@ -235,7 +240,7 @@ def get_dataset(training, image_size, config):
           validation_data=get_dataset(
               training=False, image_size=eval_size, config=config),
           validation_steps=num_eval_images // config.eval.batch_size,
-          callbacks=[ckpt_callback, tb_callback, rstr_callback],
+          callbacks=filter_callbacks([ckpt_callback, tb_callback, rstr_callback]),
           # don't log spam if running on tpus
           verbose=2 if strategy == 'tpu' else 1,
       )
@@ -245,7 +250,7 @@ def get_dataset(training, image_size, config):
             get_dataset(training=True, image_size=train_size, config=config),
             epochs=config.train.epochs,
             steps_per_epoch=steps_per_epoch,
-            callbacks=[ckpt_callback, tb_callback, rstr_callback],
+            callbacks=filter_callbacks([ckpt_callback, tb_callback, rstr_callback]),
             verbose=2 if strategy == 'tpu' else 1,
         )
       else:
@@ -274,7 +279,7 @@ def get_dataset(training, image_size, config):
               initial_epoch=start_epoch,
               epochs=end_epoch,
               steps_per_epoch=steps_per_epoch,
-              callbacks=[ckpt_callback, tb_callback, rstr_callback],
+              callbacks=filter_callbacks([ckpt_callback, tb_callback, rstr_callback]),
               verbose=2 if strategy == 'tpu' else 1,
           )
     elif FLAGS.mode == 'eval':
@@ -285,7 +290,7 @@ def get_dataset(training, image_size, config):
             get_dataset(training=False, image_size=eval_size, config=config),
             batch_size=config.eval.batch_size,
             steps=num_eval_images // config.eval.batch_size,
-            callbacks=[tb_callback, rstr_callback],
+            callbacks=filter_callbacks([tb_callback, rstr_callback]),
             verbose=2 if strategy == 'tpu' else 1,
         )
 

efficientnetv2/utils.py

@@ -336,7 +336,7 @@ def __init__(self, name, _):  # pylint: disable=super-init-not-called
 
 def scalar(name, tensor, is_tpu=True):
   """Stores a (name, Tensor) tuple in a custom collection."""
-  logging.info('Adding scale summary %s', Pair(name, tensor))
+  logging.info('Adding scalar summary %s', Pair(name, tensor))
   if is_tpu:
     tf.compat.v1.add_to_collection('scalar_summaries',
                                    Pair(name, tf.reduce_mean(tensor)))