From 31c42f87f10eccb2c621311605a745683c02d038 Mon Sep 17 00:00:00 2001 From: fuhailin Date: Fri, 19 Sep 2025 11:54:59 +0800 Subject: [PATCH 1/2] [Data] Add shuffle for DataPipeline --- deepray/datasets/datapipeline.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepray/datasets/datapipeline.py b/deepray/datasets/datapipeline.py index c0d1c676..855998c3 100644 --- a/deepray/datasets/datapipeline.py +++ b/deepray/datasets/datapipeline.py @@ -35,6 +35,7 @@ def __init__(self, context: tf.distribute.InputContext = None, **kwargs): # self.conf = Foo(flags.FLAGS.conf_file).conf self.url = None self.prebatch_size = kwargs.get("prebatch_size", None) + self.shuffle = kwargs.get("shuffle", False) @abc.abstractmethod def __len__(self): @@ -63,7 +64,7 @@ def parser(self, record): @abc.abstractmethod def build_dataset( - self, batch_size, input_file_pattern=None, is_training=True, epochs=1, shuffle=False, *args, **kwargs + self, batch_size, input_file_pattern=None, is_training=True, epochs=1, *args, **kwargs ): """ must be defined in subclass From fb5b17342284594127ae58404b8a4a61ca4e44f2 Mon Sep 17 00:00:00 2001 From: fuhailin Date: Wed, 24 Sep 2025 21:06:40 +0800 Subject: [PATCH 2/2] [Layers] Adjust var name in EV --- deepray/layers/embedding_variable.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/deepray/layers/embedding_variable.py b/deepray/layers/embedding_variable.py index 3e4f80cf..eca071e6 100644 --- a/deepray/layers/embedding_variable.py +++ b/deepray/layers/embedding_variable.py @@ -72,7 +72,7 @@ def __init__( **kwargs, ): super(EmbeddingVariable, self).__init__(name=name) - self.embedding_size = embedding_dim + self.embedding_dim = embedding_dim self.with_unique = with_unique self.world_size = get_world_size() @@ -138,21 +138,21 @@ def unique_read(self, ids, *args, **kwargs): unique_ids, idx = tf.unique(ids_flat) unique_embeddings = self.read(unique_ids) embeddings_flat = tf.gather(unique_embeddings, idx) - embeddings_shape = tf.concat([tf.shape(ids), tf.constant(self.embedding_size, shape=(1,))], 0) + embeddings_shape = tf.concat([tf.shape(ids), tf.constant(self.embedding_dim, shape=(1,))], 0) embeddings = tf.reshape(embeddings_flat, embeddings_shape) return embeddings def hvd_read(self, ids, *args, **kwargs): """ Compute embedding output for feature ids. The output shape will be (shape(ids), - embedding_size). + embedding_dim). Args: ids: feature ids of the input. It should be same dtype as the key_dtype of the layer. Returns: - A embedding output with shape (shape(ids), embedding_size). + A embedding output with shape (shape(ids), embedding_dim). """ is_ragged = isinstance(ids, tf.RaggedTensor) @@ -161,7 +161,7 @@ def hvd_read(self, ids, *args, **kwargs): ids = ids.flat_values input_shape = tf.shape(ids) - embeddings_shape = tf.concat([input_shape, [self.embedding_size]], 0) + embeddings_shape = tf.concat([input_shape, [self.embedding_dim]], 0) ids_flat = tf.reshape(ids, [-1]) @@ -179,7 +179,7 @@ def distributed_lookup(ids): lookup_result, _ = hvd.alltoall(lookup_result, splits=remote_sizes, name=f"{self.name}_alltoall_embeddings") input_shape = tf.shape(ids) - recover_shape = tf.concat((input_shape, (self.embedding_size,)), axis=0) + recover_shape = tf.concat((input_shape, (self.embedding_dim,)), axis=0) gather_indices = tf.expand_dims(tf.concat(gather_indices, axis=0), axis=-1) lookup_result = tf.scatter_nd(gather_indices, lookup_result, recover_shape) return lookup_result