[DLRM/TF2] Synthetic dataset bugfix

tgrel · tgrel · commit 15a0678fcb4d · 2022-07-26T13:00:26.000+02:00
diff --git a/TensorFlow2/Recommendation/DLRM/dataloader.py b/TensorFlow2/Recommendation/DLRM/dataloader.py
@@ -55,12 +55,14 @@ def create_input_pipelines(flags, table_ids):
         train_dataset = DummyDataset(batch_size=flags.batch_size,
                                      num_numerical_features=dataset_metadata.num_numerical_features,
                                      categorical_feature_cardinalities=local_table_sizes,
-                                     num_batches=flags.synthetic_dataset_train_batches)
+                                     num_batches=flags.synthetic_dataset_train_batches,
+                                     num_workers=hvd.size())
 
         test_dataset = DummyDataset(batch_size=flags.batch_size,
                                     num_numerical_features=dataset_metadata.num_numerical_features,
                                     categorical_feature_cardinalities=local_table_sizes,
-                                    num_batches=flags.synthetic_dataset_valid_batches)
+                                    num_batches=flags.synthetic_dataset_valid_batches,
+                                    num_workers=hvd.size())
 
     elif flags.dataset_type == 'tf_raw':
         local_categorical_feature_names = feature_spec.cat_positions_to_names(table_ids)
diff --git a/TensorFlow2/Recommendation/DLRM/datasets.py b/TensorFlow2/Recommendation/DLRM/datasets.py
@@ -28,14 +28,14 @@
 
 
 class DummyDataset:
-    def __init__(self, batch_size, num_numerical_features, categorical_feature_cardinalities, num_batches):
+    def __init__(self, batch_size, num_numerical_features, categorical_feature_cardinalities, num_batches, num_workers):
         cat_features_count = len(
             categorical_feature_cardinalities) if categorical_feature_cardinalities is not None else 0
         num_features_count = num_numerical_features if num_numerical_features is not None else 0
 
-        self.numerical_features = tf.random.uniform(shape=[batch_size, num_numerical_features], dtype=tf.float32) \
+        self.numerical_features = tf.random.uniform(shape=[batch_size // num_workers, num_numerical_features], dtype=tf.float32) \
             if num_features_count else -1
-        self.labels = tf.cast(tf.random.uniform(shape=[batch_size, 1], maxval=2, dtype=tf.int32), tf.float32)
+        self.labels = tf.cast(tf.random.uniform(shape=[batch_size // num_workers, 1], maxval=2, dtype=tf.int32), tf.float32)
         self.categorical_features = tf.concat(
             [tf.random.uniform(shape=[batch_size, 1], maxval=cardinality, dtype=tf.int32)
              for cardinality in categorical_feature_cardinalities], axis=1) if cat_features_count > 0 else -1
