Internal change

PiperOrigin-RevId: 406239748

official/projects/edgetpu/vision/tasks/image_classification_test.py

@@ -27,6 +27,19 @@ from official.projects.edgetpu.vision.configs import mobilenet_edgetpu_config
 from official.projects.edgetpu.vision.tasks import image_classification
 
 
+# Dummy ImageNet TF dataset.
+def dummy_imagenet_dataset():
+  def dummy_data(_):
+    dummy_image = tf.zeros((2, 224, 224, 3), dtype=tf.float32)
+    dummy_label = tf.zeros((2), dtype=tf.int32)
+    return (dummy_image, dummy_label)
+  dataset = tf.data.Dataset.range(1)
+  dataset = dataset.repeat()
+  dataset = dataset.map(
+      dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
+  return dataset
+
+
 class ImageClassificationTaskTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.parameters(('mobilenet_edgetpu_v2_xs'),
@@ -45,10 +58,8 @@ class ImageClassificationTaskTest(tf.test.TestCase, parameterized.TestCase):
     task = image_classification.EdgeTPUTask(config.task)
     model = task.build_model()
     metrics = task.build_metrics()
-    strategy = tf.distribute.get_strategy()
 
-    dataset = orbit.utils.make_distributed_dataset(strategy, task.build_inputs,
-                                                   config.task.train_data)
+    dataset = dummy_imagenet_dataset()
 
     iterator = iter(dataset)
     opt_factory = optimization.OptimizerFactory(config.trainer.optimizer_config)

official/projects/edgetpu/vision/tasks/semantic_segmentation_test.py

@@ -28,6 +28,25 @@ from official.projects.edgetpu.vision.tasks import semantic_segmentation as img_
 from official.vision import beta
 
 
+# Dummy ADE20K TF dataset.
+def dummy_ade20k_dataset(image_width, image_height):
+  def dummy_data(_):
+    dummy_image = tf.zeros((1, image_width, image_height, 3), dtype=tf.float32)
+    dummy_masks = tf.zeros((1, image_width, image_height, 1), dtype=tf.float32)
+    dummy_valid_masks = dummy_masks
+    dummy_image_info = tf.zeros((1, 4, 2), dtype=tf.float32)
+    return (dummy_image, {
+        'masks': dummy_masks,
+        'valid_masks': dummy_valid_masks,
+        'image_info': dummy_image_info,
+    })
+  dataset = tf.data.Dataset.range(1)
+  dataset = dataset.repeat()
+  dataset = dataset.map(
+      dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
+  return dataset
+
+
 class SemanticSegmentationTaskTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.parameters(('deeplabv3plus_mobilenet_edgetpuv2_xs_ade20k_32',),
@@ -57,10 +76,8 @@ class SemanticSegmentationTaskTest(tf.test.TestCase, parameterized.TestCase):
     task = img_seg_task.CustomSemanticSegmentationTask(config.task)
     model = task.build_model()
     metrics = task.build_metrics()
-    strategy = tf.distribute.get_strategy()
 
-    dataset = orbit.utils.make_distributed_dataset(strategy, task.build_inputs,
-                                                   config.task.train_data)
+    dataset = dummy_ade20k_dataset(32, 32)
 
     iterator = iter(dataset)
     opt_factory = optimization.OptimizerFactory(config.trainer.optimizer_config)