Internal change

PiperOrigin-RevId: 286061591

official/vision/detection/configs/retinanet_config.py

@@ -32,6 +32,7 @@ RETINANET_CFG = {
     'type': 'retinanet',
     'model_dir': '',
     'use_tpu': True,
+    'strategy_type': 'tpu',
     'train': {
         'batch_size': 64,
         'iterations_per_loop': 500,

official/vision/detection/dataloader/input_reader.py

@@ -94,7 +94,9 @@ class InputFn(object):
     dataset = dataset.cache()
 
     if self._is_training:
-      dataset = dataset.shuffle(64)
+      # Large shuffle size is critical for 2vm input pipeline. Can use small
+      # value (e.g. 64) for 1vm.
+      dataset = dataset.shuffle(1000)
     if self._num_examples > 0:
       dataset = dataset.take(self._num_examples)