[feat] optimize plan batch_size for better sharding plan when use sam…

…pler (#8)

tzrec/datasets/dataset.py

@@ -346,6 +346,14 @@ def _build_batch(self, input_data: Dict[str, pa.Array]) -> Batch:
             batch = self._data_parser.to_batch(output_data)
         return batch
 
+    @property
+    def sampled_batch_size(self) -> int:
+        """Batch size with sampler."""
+        if self._sampler:
+            return self._batch_size + self._sampler.estimated_sample_num
+        else:
+            return self._batch_size
+
 
 class BaseReader(metaclass=_reader_meta_cls):
     """Reader base class.

tzrec/datasets/sampler.py

@@ -302,6 +302,11 @@ def _parse_sparse_nodes(
         # pyre-ignore [16]
         return features, nodes.indices
 
+    @property
+    def estimated_sample_num(self) -> int:
+        """Max number of sampled num examples."""
+        raise NotImplementedError
+
 
 class NegativeSampler(BaseSampler):
     """Negative Sampler.
@@ -363,6 +368,11 @@ def get(self, input_data: Dict[str, pa.Array]) -> Dict[str, pa.Array]:
         result_dict = dict(zip(self._attr_names, features))
         return result_dict
 
+    @property
+    def estimated_sample_num(self) -> int:
+        """Estimated number of sampled num examples."""
+        return self._num_sample
+
 
 class NegativeSamplerV2(BaseSampler):
     """Negative Sampler V2.
@@ -453,6 +463,11 @@ def get(self, input_data: Dict[str, pa.Array]) -> Dict[str, pa.Array]:
         result_dict = dict(zip(self._attr_names, features))
         return result_dict
 
+    @property
+    def estimated_sample_num(self) -> int:
+        """Estimated number of sampled num examples."""
+        return self._num_sample
+
 
 class HardNegativeSampler(BaseSampler):
     """HardNegativeSampler.
@@ -548,6 +563,11 @@ def get(self, input_data: Dict[str, pa.Array]) -> Dict[str, pa.Array]:
         result_dict["hard_neg_indices"] = pa.array(hard_neg_indices)
         return result_dict
 
+    @property
+    def estimated_sample_num(self) -> int:
+        """Estimated number of sampled num examples."""
+        return self._num_sample + min(self._num_hard_sample, 8) * self._batch_size
+
 
 class HardNegativeSamplerV2(BaseSampler):
     """HardNegativeSampler.
@@ -649,6 +669,11 @@ def get(self, input_data: Dict[str, pa.Array]) -> Dict[str, pa.Array]:
         result_dict["hard_neg_indices"] = pa.array(hard_neg_indices)
         return result_dict
 
+    @property
+    def estimated_sample_num(self) -> int:
+        """Estimated number of sampled num examples."""
+        return self._num_sample + min(self._num_hard_sample, 8) * self._batch_size
+
 
 class TDMSampler(BaseSampler):
     """TDM training sampler.
@@ -840,6 +865,13 @@ def get(self, input_data: Dict[str, pa.Array]) -> Dict[str, pa.Array]:
 
         return pos_result_dict, neg_result_dict
 
+    @property
+    def estimated_sample_num(self) -> int:
+        """Estimated number of sampled num examples."""
+        return (
+            sum(self._layer_num_sample) + len(self._layer_num_sample) - 2
+        ) * self._batch_size
+
 
 class TDMPredictSampler(BaseSampler):
     """TDM predict sampler.
@@ -913,3 +945,8 @@ def get(self, input_ids: pa.Array) -> Dict[str, pa.Array]:
         pos_result_dict = dict(zip(self._attr_names[1:], pos_fea_result))
 
         return pos_result_dict
+
+    @property
+    def estimated_sample_num(self) -> int:
+        """Estimated number of sampled num examples."""
+        return min((2 ** (self._max_level - 1)), 800) * self._batch_size

tzrec/datasets/sampler_test.py

@@ -158,6 +158,7 @@ def _sampler_worker(res):
                 ],
                 batch_size=4,
             )
+            assert sampler.estimated_sample_num == 8
             sampler.init_cluster()
             sampler.launch_server()
             sampler.init()
@@ -285,6 +286,7 @@ def _sampler_worker(res):
                 ],
                 batch_size=4,
             )
+            assert sampler.estimated_sample_num == 8
             sampler.init_cluster()
             sampler.launch_server()
             sampler.init()
@@ -332,6 +334,7 @@ def _sampler_worker(res):
                 ],
                 batch_size=4,
             )
+            assert sampler.estimated_sample_num == 40
             sampler.init_cluster()
             sampler.launch_server()
             sampler.init()
@@ -381,6 +384,7 @@ def _sampler_worker(res):
                 ],
                 batch_size=4,
             )
+            assert sampler.estimated_sample_num == 40
             sampler.init_cluster()
             sampler.launch_server()
             sampler.init()
@@ -426,6 +430,7 @@ def _sampler_worker(pos_res, neg_res):
                 ],
                 batch_size=4,
             )
+            assert sampler.estimated_sample_num == 76
             sampler.init_cluster()
             sampler.launch_server()
             sampler.init()
@@ -487,6 +492,7 @@ def _sampler_worker(res):
                 ],
                 batch_size=4,
             )
+            assert sampler.estimated_sample_num == 128
             sampler.init_cluster()
             sampler.launch_server()
             sampler.init()

tzrec/main.py

@@ -544,7 +544,8 @@ def train_and_evaluate(
 
     planner = create_planner(
         device=device,
-        batch_size=data_config.batch_size,
+        # pyre-ignore [16]
+        batch_size=train_dataloader.dataset.sampled_batch_size,
     )
     plan = planner.collective_plan(
         model, get_default_sharders(), dist.GroupMember.WORLD
@@ -662,7 +663,8 @@ def evaluate(
 
     planner = create_planner(
         device=device,
-        batch_size=data_config.batch_size,
+        # pyre-ignore [16]
+        batch_size=eval_dataloader.dataset.sampled_batch_size,
     )
     plan = planner.collective_plan(
         model, get_default_sharders(), dist.GroupMember.WORLD