title: How to speed up tf.data.Dataset.from_generator()
body: <p>In tensorflow2.0, I want to train a skip-gram model with nce loss. tf.data.Dataset.from_tensor_slices() is not suitable because the input file is really huge. So I write a dataset generator class like this:</p>
  
  <pre class="lang-py prettyprint-override"><code>class DataSet:
      """"""
  
      def __init__(self, args, vocab):
          self.args = args
          self.vocab = vocab
  
      def generator(self):
          """a generator function, it will return skip-gram sample or cbow sample"""
          with open(self.args.input) as f_input:
              for line in tqdm.tqdm(f_input.readlines()):
                  tokens = line.strip().split()
                  tokens_indices = self.vocab.indices(tokens)
                  for index, target_word in enumerate(tokens_indices):
                      context_words = list()
                      begin = index - self.args.window_size if index - self.args.window_size &gt; 0 else 0
                      end = index + 1 + self.args.window_size if index + self.args.window_size + 1 &lt; len(tokens_indices) else len(
                          tokens_indices)
                      context_words.extend(tokens_indices[begin:index])
                      context_words.extend(tokens_indices[index + 1:end])
                      if self.args.cbow &gt; 0:
                          yield context_words, target_word
                      else:
                          for i in range(len(context_words)):
                              yield target_word, context_words[i]
  
      def dataset(self):
          """Using tf.data.Dataset.from_generator() to return sample"""
          if self.args.cbow:
              dataset = tf.data.Dataset.from_generator(
                  self.generator,
                  (tf.int32, tf.int32),
                  (tf.TensorShape([None]), tf.TensorShape([]))
              )
          else:
              dataset = tf.data.Dataset.from_generator(
                  self.generator,
                  (tf.int32, tf.int32),
                  (tf.TensorShape([]), tf.TensorShape([]))
              )
  
          return dataset
  </code></pre>
  
  <p>Then I test my code with follow:</p>
  
  <pre class="lang-py prettyprint-override"><code>dataset = DataSet(args, vocab).dataset()
  iterator = dataset.make_one_shot_iterator()
  for batch, (x,y) in enumerate(dataset.batch(128)):
      pass
  print(batch, x.shape, y.shape)
  </code></pre>
  
  <p>But it cost a lot of time to iterate all lines(about 10 minutes / 15000 lines in MacBook pro 2012). Does there any methods can speed up the code?</p>