Fix errors in classifier trainer when there is no distribution strategy

PiperOrigin-RevId: 306495321

Fix errors in classifier trainer when there is no distribution strategy
PiperOrigin-RevId: 306495321
02cf1f87 · Zongwei Zhou · A. Unique TensorFlower · 0bfdad07 · 02cf1f87
Commit 02cf1f87 authored Apr 14, 2020 by Zongwei Zhou Committed by A. Unique TensorFlower Apr 14, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 2 deletions

official/vision/image_classification/classifier_trainer.py official/vision/image_classification/classifier_trainer.py +4 -2

No files found.
--- a/official/vision/image_classification/classifier_trainer.py
+++ b/official/vision/image_classification/classifier_trainer.py
@@ -109,7 +109,8 @@ def _get_dataset_builders(params: base_configs.ExperimentConfig,
    logging.warning('label_smoothing not applied, so datasets will not be one '
                    'hot encoded.')

-  num_devices = strategy.num_replicas_in_sync
+  num_devices = strategy.num_replicas_in_sync if strategy else 1
+
  image_size = get_image_size_from_model(params)

  dataset_configs = [
@@ -309,7 +310,8 @@ def train_and_eval(

  strategy_scope = distribution_utils.get_strategy_scope(strategy)

-  logging.info('Detected %d devices.', strategy.num_replicas_in_sync)
+  logging.info('Detected %d devices.',
+               strategy.num_replicas_in_sync if strategy else 1)

  label_smoothing = params.model.loss.label_smoothing
  one_hot = label_smoothing and label_smoothing > 0