put shard first
@@ -183,15 +183,25 @@ class ParallelTextDataLayer(DataLayer):
[SpecialTextTokens.EOS_ID.value], self._pad_lengths_to_eight), dtype="int32")
def build_graph(self):
- _sources = tf.data.TextLineDataset(self.source_file)\
- .map(lambda line: tf.py_func(func=self._src_token_to_id, inp=[line],
+
+ _sources = tf.data.TextLineDataset(self.source_file)
+ _targets = tf.data.TextLineDataset(self.target_file)
+
+ if self._num_workers > 1:
+ #_src_tgt_dataset = _src_tgt_dataset\
+ # .shard(num_shards=self._num_workers, index=self._worker_id)
+ _sources = _sources.shard(num_shards=self._num_workers,
+ index=self._worker_id)
+ _targets = _targets.shard(num_shards=self._num_workers,
+ index=self._worker_id)
+
+ _sources = _sources.map(lambda line: tf.py_func(func=self._src_token_to_id, inp=[line],
Tout=[tf.int32], stateful=False),
num_parallel_calls=self._map_parallel_calls) \
.map(lambda tokens: (tokens, tf.size(tokens)),
num_parallel_calls=self._map_parallel_calls)
- _targets = tf.data.TextLineDataset(self.target_file) \
- .map(lambda line: tf.py_func(func=self._tgt_token_to_id, inp=[line],
+ _targets = _targets.map(lambda line: tf.py_func(func=self._tgt_token_to_id, inp=[line],
Tout=[tf.int32], stateful=False),
num_parallel_calls=self._map_parallel_calls) \
.map(lambda tokens: (tokens, tf.size(tokens)),
@@ -202,10 +212,6 @@ class ParallelTextDataLayer(DataLayer):
tf.less_equal(t2[1], self.max_len))
).cache()
- if self._num_workers > 1:
- _src_tgt_dataset = _src_tgt_dataset\
- .shard(num_shards=self._num_workers, index=self._worker_id)
-
if self.params['shuffle']:
bf_size = self.get_size_in_samples() if self._shuffle_buffer_size == -1 \
else self._shuffle_buffer_size
