Rm num_shards from Problem. Problems specify sharding themselves.

PiperOrigin-RevId: 163281576

Author: Ryan Sepassi

README.md

@@ -86,7 +86,6 @@ mkdir -p $DATA_DIR $TMP_DIR $TRAIN_DIR
 t2t-datagen \
   --data_dir=$DATA_DIR \
   --tmp_dir=$TMP_DIR \
-  --num_shards=100 \
   --problem=$PROBLEM
 
 # Train

tensor2tensor/bin/t2t-datagen

@@ -63,7 +63,8 @@ flags.DEFINE_string("problem", "",
                     "The name of the problem to generate data for.")
 flags.DEFINE_string("exclude_problems", "",
                     "Comma-separates list of problems to exclude.")
-flags.DEFINE_integer("num_shards", 10, "How many shards to use.")
+flags.DEFINE_integer("num_shards", 0, "How many shards to use. Ignored for "
+                     "registered Problems.")
 flags.DEFINE_integer("max_cases", 0,
                      "Maximum number of cases to generate (unbounded if 0).")
 flags.DEFINE_integer("random_seed", 429459, "Random seed to use.")
@@ -252,7 +253,7 @@ def generate_data_for_problem(problem):
   if isinstance(dev_gen, int):
     # The dev set and test sets are generated as extra shards using the
     # training generator.  The integer specifies the number of training
-    # shards.  FLAGS.num_shards is ignored.
+    # shards. FLAGS.num_shards is ignored.
     num_training_shards = dev_gen
     tf.logging.info("Generating data for %s.", problem)
     all_output_files = generator_utils.combined_data_filenames(
@@ -263,10 +264,11 @@ def generate_data_for_problem(problem):
   else:
     # usual case - train data and dev data are generated using separate
     # generators.
+    num_shards = FLAGS.num_shards or 10
     tf.logging.info("Generating training data for %s.", problem)
     train_output_files = generator_utils.train_data_filenames(
         problem + generator_utils.UNSHUFFLED_SUFFIX, FLAGS.data_dir,
-        FLAGS.num_shards)
+        num_shards)
     generator_utils.generate_files(training_gen(), train_output_files,
                                    FLAGS.max_cases)
     tf.logging.info("Generating development data for %s.", problem)
@@ -282,11 +284,12 @@ def generate_data_for_problem(problem):
 
 def generate_data_for_registered_problem(problem_name):
   tf.logging.info("Generating training data for %s.", problem_name)
+  if FLAGS.num_shards:
+    raise ValueError("--num_shards should not be set for registered Problem.")
   problem = registry.problem(problem_name)
   task_id = None if FLAGS.task_id < 0 else FLAGS.task_id
   problem.generate_data(os.path.expanduser(FLAGS.data_dir),
                         os.path.expanduser(FLAGS.tmp_dir),
-                        num_shards=FLAGS.num_shards,
                         task_id=task_id)
 
 

tensor2tensor/data_generators/algorithmic.py

@@ -66,10 +66,7 @@ def dev_size(self):
   def num_shards(self):
     return 10
 
-  def generate_data(self, data_dir, _, num_shards=None, task_id=-1):
-    if num_shards is None:
-      num_shards = self.num_shards
-
+  def generate_data(self, data_dir, _, task_id=-1):
     def generator_eos(generator):
       """Shift by NUM_RESERVED_IDS and append EOS token."""
       for case in generator:
@@ -87,7 +84,7 @@ def generator_eos(generator):
 
     utils.generate_dataset_and_shuffle(
         train_generator_eos(),
-        self.training_filepaths(data_dir, num_shards, shuffled=True),
+        self.training_filepaths(data_dir, self.num_shards, shuffled=True),
         dev_generator_eos(),
         self.dev_filepaths(data_dir, 1, shuffled=True),
         shuffle=False)
@@ -254,7 +251,7 @@ def zipf_distribution(nbr_symbols, alpha):
 
 
 def zipf_random_sample(distr_map, sample_len):
-  """Helper function: Generate a random Zipf sample of given lenght.
+  """Helper function: Generate a random Zipf sample of given length.
 
   Args:
     distr_map: list of float, Zipf's distribution over nbr_symbols.
@@ -287,7 +284,7 @@ def reverse_generator_nlplike(nbr_symbols,
     max_length: integer, maximum length of sequences to generate.
     nbr_cases: the number of cases to generate.
     scale_std_dev: float, Normal distribution's standard deviation scale factor
-      used to draw the lenght of sequence. Default = 1% of the max_length.
+      used to draw the length of sequence. Default = 1% of the max_length.
     alpha: float, Zipf's Law Distribution parameter. Default = 1.5.
       Usually for modelling natural text distribution is in
       the range [1.1-1.6].

tensor2tensor/data_generators/genetics.py

@@ -87,10 +87,11 @@ def feature_encoders(self, data_dir):
         "targets": text_encoder.TextEncoder()
     }
 
-  def generate_data(self, data_dir, tmp_dir, num_shards=None, task_id=-1):
-    if num_shards is None:
-      num_shards = 100
+  @property
+  def num_shards(self):
+    return 100
 
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
     try:
       # Download source data if download_url specified
       h5_filepath = generator_utils.maybe_download(tmp_dir, self.h5_file,
@@ -109,7 +110,7 @@ def generate_data(self, data_dir, tmp_dir, num_shards=None, task_id=-1):
     # Collect created shard processes to start and join
     processes = []
 
-    datasets = [(self.training_filepaths, num_shards, "train",
+    datasets = [(self.training_filepaths, self.num_shards, "train",
                  num_train_examples), (self.dev_filepaths, 1, "valid",
                                        num_dev_examples),
                 (self.test_filepaths, 1, "test", num_test_examples)]
@@ -124,9 +125,10 @@ def generate_data(self, data_dir, tmp_dir, num_shards=None, task_id=-1):
                   start_idx, end_idx))
         processes.append(p)
 
-    # Start and wait for processes in batches
-    assert len(processes) == num_shards + 2  # 1 per training shard + dev + test
+    # 1 per training shard + dev + test
+    assert len(processes) == self.num_shards + 2
 
+    # Start and wait for processes in batches
     num_batches = int(
         math.ceil(float(len(processes)) / MAX_CONCURRENT_PROCESSES))
     for i in xrange(num_batches):

tensor2tensor/data_generators/image.py

@@ -338,7 +338,7 @@ def example_reading_spec(self, label_key=None):
 class ImageFSNS(ImageProblem):
   """Problem spec for French Street Name recognition."""
 
-  def generate_data(self, data_dir, tmp_dir, num_shards=None, task_id=-1):
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
     list_url = ("https://raw.githubusercontent.com/tensorflow/models/master/"
                 "street/python/fsns_urls.txt")
     fsns_urls = generator_utils.maybe_download(

tensor2tensor/data_generators/problem.py

@@ -135,7 +135,7 @@ class Problem(object):
   # BEGIN SUBCLASS INTERFACE
   # ============================================================================
 
-  def generate_data(self, data_dir, tmp_dir, num_shards=None, task_id=-1):
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
     raise NotImplementedError()
 
   def hparams(self, defaults, model_hparams):

tensor2tensor/data_generators/wmt.py

@@ -83,12 +83,10 @@ def vocab_name(self):
   def vocab_file(self):
     return "%s.%d" % (self.vocab_name, self.targeted_vocab_size)
 
-  def generate_data(self, data_dir, tmp_dir, num_shards=None, task_id=-1):
-    if num_shards is None:
-      num_shards = self.num_shards
+  def generate_data(self, data_dir, tmp_dir, task_id=-1):
     generator_utils.generate_dataset_and_shuffle(
         self.train_generator(data_dir, tmp_dir, True),
-        self.training_filepaths(data_dir, num_shards, shuffled=False),
+        self.training_filepaths(data_dir, self.num_shards, shuffled=False),
         self.dev_generator(data_dir, tmp_dir),
         self.dev_filepaths(data_dir, 1, shuffled=False))