Internal change

PiperOrigin-RevId: 351691859

official/vision/beta/configs/image_classification.py

@@ -35,6 +35,7 @@ class DataConfig(cfg.DataConfig):
   shuffle_buffer_size: int = 10000
   cycle_length: int = 10
   aug_policy: Optional[str] = None  # None, 'autoaug', or 'randaug'
+  file_type: str = 'tfrecord'  # tfrecord, or sstable
 
 
 @dataclasses.dataclass

official/vision/beta/tasks/image_classification.py

@@ -22,6 +22,7 @@ from official.core import task_factory
 from official.modeling import tf_utils
 from official.vision.beta.configs import image_classification as exp_cfg
 from official.vision.beta.dataloaders import classification_input
+from official.vision.beta.dataloaders import dataset_fn
 from official.vision.beta.modeling import factory
 
 
@@ -86,7 +87,7 @@ class ImageClassificationTask(base_task.Task):
 
     reader = input_reader.InputReader(
         params,
-        dataset_fn=tf.data.TFRecordDataset,
+        dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
         decoder_fn=decoder.decode,
         parser_fn=parser.parse_fn(params.is_training))