Improve usability of batch_size in as_dataset function.

Author: ChanchalKumarMaji

tensorflow_datasets/core/dataset_builder.py

@@ -308,7 +308,7 @@ def download_and_prepare(self, download_dir=None, download_config=None):
   @api_utils.disallow_positional_args
   def as_dataset(self,
                  split=None,
-                 batch_size=1,
+                 batch_size=None,
                  shuffle_files=None,
                  as_supervised=False):
     """Constructs a `tf.data.Dataset`.
@@ -376,7 +376,7 @@ def _build_single_dataset(self, split, shuffle_files, batch_size,
       batch_size = self.info.splits.total_num_examples or sys.maxsize
 
     dataset = self._as_dataset(split=split, shuffle_files=shuffle_files)
-    if batch_size > 1:
+    if batch_size:
       # Use padded_batch so that features with unknown shape are supported.
       padded_shapes = self.info.features.shape
       dataset = dataset.padded_batch(batch_size, padded_shapes)

tensorflow_datasets/core/dataset_builder_test.py

@@ -423,6 +423,16 @@ def test_with_batch_size(self):
     self.assertEqual(10, x2.shape[0])
     self.assertEqual(10, x3.shape[0])
     self.assertEqual(sum(range(30)), int(x1.sum() + x2.sum() + x3.sum()))
+    
+    # By default batch_size is None and won't add a batch dimension
+    ds = self.builder.as_dataset(split=splits_lib.Split.TRAIN)
+    self.assertEqual(0, len(ds.output_shapes["x"]))
+    # Setting batch_size=1 will add an extra batch dimension
+    ds = self.builder.as_dataset(split=splits_lib.Split.TRAIN, batch_size=1)
+    self.assertEqual(1, len(ds.output_shapes["x"]))
+    # Setting batch_size=2 will add an extra batch dimension
+    ds = self.builder.as_dataset(split=splits_lib.Split.TRAIN, batch_size=2)
+    self.assertEqual(1, len(ds.output_shapes["x"]))
 
   @testing.run_in_graph_and_eager_modes()
   def test_supervised_keys(self):

tensorflow_datasets/core/registered.py

@@ -176,7 +176,7 @@ def builder(name, **builder_init_kwargs):
 def load(name,
          split=None,
          data_dir=None,
-         batch_size=1,
+         batch_size=None,
          download=True,
          as_supervised=False,
          with_info=False,

tensorflow_datasets/core/registered_test.py

@@ -120,7 +120,7 @@ def test_load(self):
     self.assertFalse(builder.download_called)
     self.assertEqual(splits.Split.TEST,
                      builder.as_dataset_kwargs.pop("split"))
-    self.assertEqual(1, builder.as_dataset_kwargs.pop("batch_size"))
+    self.assertEqual(None, builder.as_dataset_kwargs.pop("batch_size"))
     self.assertFalse(builder.as_dataset_kwargs.pop("as_supervised"))
     self.assertEqual(builder.as_dataset_kwargs, as_dataset_kwargs)
     self.assertEqual(dict(data_dir=data_dir, k1=1), builder.kwargs)
@@ -130,6 +130,22 @@ def test_load(self):
         download=True, as_dataset_kwargs=as_dataset_kwargs)
     self.assertTrue(builder.as_dataset_called)
     self.assertTrue(builder.download_called)
+    
+    # Tests for different batch_size
+    # By default batch_size=None
+    builder = registered.load(
+        name=name, split=splits.Split.TEST, data_dir=data_dir)
+    self.assertEqual(None, builder.as_dataset_kwargs.pop("batch_size"))
+    # Setting batch_size=1
+    builder = registered.load(
+        name=name, split=splits.Split.TEST, data_dir=data_dir,
+        batch_size=1)
+    self.assertEqual(1, builder.as_dataset_kwargs.pop("batch_size"))
+    # Setting batch_size=2
+    builder = registered.load(
+        name=name, split=splits.Split.TEST, data_dir=data_dir,
+        batch_size=2)
+    self.assertEqual(2, builder.as_dataset_kwargs.pop("batch_size"))
 
   def test_load_all_splits(self):
     name = "empty_dataset_builder"