We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Error training latest version of MusicVAE
#### Steps to reproduce the issue

- Create new magenta environment as well as cloning magenta repo:
`curl https://raw.githubusercontent.com/tensorflow/magenta/master/magenta/tools/magenta-install.sh > /tmp/magenta-install.sh`
`bash /tmp/magenta-install.sh`
`git clone https://github.com/tensorflow/magenta.git`

- Training simple model:
`music_vae_train --config=cat-mel_2bar_small --run_dir=/tmp/cat-mel_2bar_small/ --mode=train --examples_path=data/notesequences.tfrecord`

#### Expected results:

- Training loss and global_step print on console.

#### Actual output:

- `line 123, in _get_input_tensors
    iterator = dataset.make_one_shot_iterator()
AttributeError: 'PrefetchDataset' object has no attribute 'make_one_shot_iterator'`

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `magenta/models/music_vae/music_vae_train.py`
Content:
```
1 # Copyright 2020 The Magenta Authors.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 # Lint as: python3
16 """MusicVAE training script."""
17 
18 from __future__ import absolute_import
19 from __future__ import division
20 from __future__ import print_function
21 
22 import os
23 
24 from magenta.models.music_vae import configs
25 from magenta.models.music_vae import data
26 import tensorflow.compat.v1 as tf
27 from tensorflow.contrib import training as contrib_training
28 
29 flags = tf.app.flags
30 FLAGS = flags.FLAGS
31 
32 flags.DEFINE_string(
33     'master', '',
34     'The TensorFlow master to use.')
35 flags.DEFINE_string(
36     'examples_path', None,
37     'Path to a TFRecord file of NoteSequence examples. Overrides the config.')
38 flags.DEFINE_string(
39     'tfds_name', None,
40     'TensorFlow Datasets dataset name to use. Overrides the config.')
41 flags.DEFINE_string(
42     'run_dir', None,
43     'Path where checkpoints and summary events will be located during '
44     'training and evaluation. Separate subdirectories `train` and `eval` '
45     'will be created within this directory.')
46 flags.DEFINE_integer(
47     'num_steps', 200000,
48     'Number of training steps or `None` for infinite.')
49 flags.DEFINE_integer(
50     'eval_num_batches', None,
51     'Number of batches to use during evaluation or `None` for all batches '
52     'in the data source.')
53 flags.DEFINE_integer(
54     'checkpoints_to_keep', 100,
55     'Maximum number of checkpoints to keep in `train` mode or 0 for infinite.')
56 flags.DEFINE_integer(
57     'keep_checkpoint_every_n_hours', 1,
58     'In addition to checkpoints_to_keep, keep a checkpoint every N hours.')
59 flags.DEFINE_string(
60     'mode', 'train',
61     'Which mode to use (`train` or `eval`).')
62 flags.DEFINE_string(
63     'config', '',
64     'The name of the config to use.')
65 flags.DEFINE_string(
66     'hparams', '',
67     'A comma-separated list of `name=value` hyperparameter values to merge '
68     'with those in the config.')
69 flags.DEFINE_bool(
70     'cache_dataset', True,
71     'Whether to cache the dataset in memory for improved training speed. May '
72     'cause memory errors for very large datasets.')
73 flags.DEFINE_integer(
74     'task', 0,
75     'The task number for this worker.')
76 flags.DEFINE_integer(
77     'num_ps_tasks', 0,
78     'The number of parameter server tasks.')
79 flags.DEFINE_integer(
80     'num_sync_workers', 0,
81     'The number of synchronized workers.')
82 flags.DEFINE_integer(
83     'num_data_threads', 4,
84     'The number of data preprocessing threads.')
85 flags.DEFINE_string(
86     'eval_dir_suffix', '',
87     'Suffix to add to eval output directory.')
88 flags.DEFINE_string(
89     'log', 'INFO',
90     'The threshold for what messages will be logged: '
91     'DEBUG, INFO, WARN, ERROR, or FATAL.')
92 
93 
94 # Should not be called from within the graph to avoid redundant summaries.
95 def _trial_summary(hparams, examples_path, output_dir):
96   """Writes a tensorboard text summary of the trial."""
97 
98   examples_path_summary = tf.summary.text(
99       'examples_path', tf.constant(examples_path, name='examples_path'),
100       collections=[])
101 
102   hparams_dict = hparams.values()
103 
104   # Create a markdown table from hparams.
105   header = '| Key | Value |\n| :--- | :--- |\n'
106   keys = sorted(hparams_dict.keys())
107   lines = ['| %s | %s |' % (key, str(hparams_dict[key])) for key in keys]
108   hparams_table = header + '\n'.join(lines) + '\n'
109 
110   hparam_summary = tf.summary.text(
111       'hparams', tf.constant(hparams_table, name='hparams'), collections=[])
112 
113   with tf.Session() as sess:
114     writer = tf.summary.FileWriter(output_dir, graph=sess.graph)
115     writer.add_summary(examples_path_summary.eval())
116     writer.add_summary(hparam_summary.eval())
117     writer.close()
118 
119 
120 def _get_input_tensors(dataset, config):
121   """Get input tensors from dataset."""
122   batch_size = config.hparams.batch_size
123   iterator = dataset.make_one_shot_iterator()
124   (input_sequence, output_sequence, control_sequence,
125    sequence_length) = iterator.get_next()
126   input_sequence.set_shape(
127       [batch_size, None, config.data_converter.input_depth])
128   output_sequence.set_shape(
129       [batch_size, None, config.data_converter.output_depth])
130   if not config.data_converter.control_depth:
131     control_sequence = None
132   else:
133     control_sequence.set_shape(
134         [batch_size, None, config.data_converter.control_depth])
135   sequence_length.set_shape([batch_size] + sequence_length.shape[1:].as_list())
136 
137   return {
138       'input_sequence': input_sequence,
139       'output_sequence': output_sequence,
140       'control_sequence': control_sequence,
141       'sequence_length': sequence_length
142   }
143 
144 
145 def train(train_dir,
146           config,
147           dataset_fn,
148           checkpoints_to_keep=5,
149           keep_checkpoint_every_n_hours=1,
150           num_steps=None,
151           master='',
152           num_sync_workers=0,
153           num_ps_tasks=0,
154           task=0):
155   """Train loop."""
156   tf.gfile.MakeDirs(train_dir)
157   is_chief = (task == 0)
158   if is_chief:
159     _trial_summary(
160         config.hparams, config.train_examples_path or config.tfds_name,
161         train_dir)
162   with tf.Graph().as_default():
163     with tf.device(tf.train.replica_device_setter(
164         num_ps_tasks, merge_devices=True)):
165 
166       model = config.model
167       model.build(config.hparams,
168                   config.data_converter.output_depth,
169                   is_training=True)
170 
171       optimizer = model.train(**_get_input_tensors(dataset_fn(), config))
172 
173       hooks = []
174       if num_sync_workers:
175         optimizer = tf.train.SyncReplicasOptimizer(
176             optimizer,
177             num_sync_workers)
178         hooks.append(optimizer.make_session_run_hook(is_chief))
179 
180       grads, var_list = list(zip(*optimizer.compute_gradients(model.loss)))
181       global_norm = tf.global_norm(grads)
182       tf.summary.scalar('global_norm', global_norm)
183 
184       if config.hparams.clip_mode == 'value':
185         g = config.hparams.grad_clip
186         clipped_grads = [tf.clip_by_value(grad, -g, g) for grad in grads]
187       elif config.hparams.clip_mode == 'global_norm':
188         clipped_grads = tf.cond(
189             global_norm < config.hparams.grad_norm_clip_to_zero,
190             lambda: tf.clip_by_global_norm(  # pylint:disable=g-long-lambda
191                 grads, config.hparams.grad_clip, use_norm=global_norm)[0],
192             lambda: [tf.zeros(tf.shape(g)) for g in grads])
193       else:
194         raise ValueError(
195             'Unknown clip_mode: {}'.format(config.hparams.clip_mode))
196       train_op = optimizer.apply_gradients(
197           list(zip(clipped_grads, var_list)),
198           global_step=model.global_step,
199           name='train_step')
200 
201       logging_dict = {'global_step': model.global_step,
202                       'loss': model.loss}
203 
204       hooks.append(tf.train.LoggingTensorHook(logging_dict, every_n_iter=100))
205       if num_steps:
206         hooks.append(tf.train.StopAtStepHook(last_step=num_steps))
207 
208       scaffold = tf.train.Scaffold(
209           saver=tf.train.Saver(
210               max_to_keep=checkpoints_to_keep,
211               keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours))
212       contrib_training.train(
213           train_op=train_op,
214           logdir=train_dir,
215           scaffold=scaffold,
216           hooks=hooks,
217           save_checkpoint_secs=60,
218           master=master,
219           is_chief=is_chief)
220 
221 
222 def evaluate(train_dir,
223              eval_dir,
224              config,
225              dataset_fn,
226              num_batches,
227              master=''):
228   """Evaluate the model repeatedly."""
229   tf.gfile.MakeDirs(eval_dir)
230 
231   _trial_summary(
232       config.hparams, config.eval_examples_path or config.tfds_name, eval_dir)
233   with tf.Graph().as_default():
234     model = config.model
235     model.build(config.hparams,
236                 config.data_converter.output_depth,
237                 is_training=False)
238 
239     eval_op = model.eval(
240         **_get_input_tensors(dataset_fn().take(num_batches), config))
241 
242     hooks = [
243         contrib_training.StopAfterNEvalsHook(num_batches),
244         contrib_training.SummaryAtEndHook(eval_dir)
245     ]
246     contrib_training.evaluate_repeatedly(
247         train_dir,
248         eval_ops=eval_op,
249         hooks=hooks,
250         eval_interval_secs=60,
251         master=master)
252 
253 
254 def run(config_map,
255         tf_file_reader=tf.data.TFRecordDataset,
256         file_reader=tf.python_io.tf_record_iterator):
257   """Load model params, save config file and start trainer.
258 
259   Args:
260     config_map: Dictionary mapping configuration name to Config object.
261     tf_file_reader: The tf.data.Dataset class to use for reading files.
262     file_reader: The Python reader to use for reading files.
263 
264   Raises:
265     ValueError: if required flags are missing or invalid.
266   """
267   if not FLAGS.run_dir:
268     raise ValueError('Invalid run directory: %s' % FLAGS.run_dir)
269   run_dir = os.path.expanduser(FLAGS.run_dir)
270   train_dir = os.path.join(run_dir, 'train')
271 
272   if FLAGS.mode not in ['train', 'eval']:
273     raise ValueError('Invalid mode: %s' % FLAGS.mode)
274 
275   if FLAGS.config not in config_map:
276     raise ValueError('Invalid config: %s' % FLAGS.config)
277   config = config_map[FLAGS.config]
278   if FLAGS.hparams:
279     config.hparams.parse(FLAGS.hparams)
280   config_update_map = {}
281   if FLAGS.examples_path:
282     config_update_map['%s_examples_path' % FLAGS.mode] = os.path.expanduser(
283         FLAGS.examples_path)
284   if FLAGS.tfds_name:
285     if FLAGS.examples_path:
286       raise ValueError(
287           'At most one of --examples_path and --tfds_name can be set.')
288     config_update_map['tfds_name'] = FLAGS.tfds_name
289     config_update_map['eval_examples_path'] = None
290     config_update_map['train_examples_path'] = None
291   config = configs.update_config(config, config_update_map)
292   if FLAGS.num_sync_workers:
293     config.hparams.batch_size //= FLAGS.num_sync_workers
294 
295   if FLAGS.mode == 'train':
296     is_training = True
297   elif FLAGS.mode == 'eval':
298     is_training = False
299   else:
300     raise ValueError('Invalid mode: {}'.format(FLAGS.mode))
301 
302   def dataset_fn():
303     return data.get_dataset(
304         config,
305         tf_file_reader=tf_file_reader,
306         num_threads=FLAGS.num_data_threads,
307         is_training=is_training,
308         cache_dataset=FLAGS.cache_dataset)
309 
310   if is_training:
311     train(
312         train_dir,
313         config=config,
314         dataset_fn=dataset_fn,
315         checkpoints_to_keep=FLAGS.checkpoints_to_keep,
316         keep_checkpoint_every_n_hours=FLAGS.keep_checkpoint_every_n_hours,
317         num_steps=FLAGS.num_steps,
318         master=FLAGS.master,
319         num_sync_workers=FLAGS.num_sync_workers,
320         num_ps_tasks=FLAGS.num_ps_tasks,
321         task=FLAGS.task)
322   else:
323     num_batches = FLAGS.eval_num_batches or data.count_examples(
324         config.eval_examples_path,
325         config.tfds_name,
326         config.data_converter,
327         file_reader) // config.hparams.batch_size
328     eval_dir = os.path.join(run_dir, 'eval' + FLAGS.eval_dir_suffix)
329     evaluate(
330         train_dir,
331         eval_dir,
332         config=config,
333         dataset_fn=dataset_fn,
334         num_batches=num_batches,
335         master=FLAGS.master)
336 
337 
338 def main(unused_argv):
339   tf.logging.set_verbosity(FLAGS.log)
340   run(configs.CONFIG_MAP)
341 
342 
343 def console_entry_point():
344   tf.app.run(main)
345 
346 
347 if __name__ == '__main__':
348   console_entry_point()
349
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/magenta/models/music_vae/music_vae_train.py b/magenta/models/music_vae/music_vae_train.py
--- a/magenta/models/music_vae/music_vae_train.py
+++ b/magenta/models/music_vae/music_vae_train.py
@@ -120,7 +120,7 @@
 def _get_input_tensors(dataset, config):
   """Get input tensors from dataset."""
   batch_size = config.hparams.batch_size
-  iterator = dataset.make_one_shot_iterator()
+  iterator = tf.data.make_one_shot_iterator(dataset)
   (input_sequence, output_sequence, control_sequence,
    sequence_length) = iterator.get_next()
   input_sequence.set_shape(

{"golden_diff": "diff --git a/magenta/models/music_vae/music_vae_train.py b/magenta/models/music_vae/music_vae_train.py\n--- a/magenta/models/music_vae/music_vae_train.py\n+++ b/magenta/models/music_vae/music_vae_train.py\n@@ -120,7 +120,7 @@\n def _get_input_tensors(dataset, config):\n   \"\"\"Get input tensors from dataset.\"\"\"\n   batch_size = config.hparams.batch_size\n-  iterator = dataset.make_one_shot_iterator()\n+  iterator = tf.data.make_one_shot_iterator(dataset)\n   (input_sequence, output_sequence, control_sequence,\n    sequence_length) = iterator.get_next()\n   input_sequence.set_shape(\n", "issue": "Error training latest version of MusicVAE\n#### Steps to reproduce the issue\r\n\r\n- Create new magenta environment as well as cloning magenta repo:\r\n`curl https://raw.githubusercontent.com/tensorflow/magenta/master/magenta/tools/magenta-install.sh > /tmp/magenta-install.sh`\r\n`bash /tmp/magenta-install.sh`\r\n`git clone https://github.com/tensorflow/magenta.git`\r\n\r\n- Training simple model:\r\n`music_vae_train --config=cat-mel_2bar_small --run_dir=/tmp/cat-mel_2bar_small/ --mode=train --examples_path=data/notesequences.tfrecord`\r\n\r\n#### Expected results:\r\n\r\n- Training loss and global_step print on console.\r\n\r\n#### Actual output:\r\n\r\n- `line 123, in _get_input_tensors\r\n    iterator = dataset.make_one_shot_iterator()\r\nAttributeError: 'PrefetchDataset' object has no attribute 'make_one_shot_iterator'`\n", "before_files": [{"content": "# Copyright 2020 The Magenta Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Lint as: python3\n\"\"\"MusicVAE training script.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport os\n\nfrom magenta.models.music_vae import configs\nfrom magenta.models.music_vae import data\nimport tensorflow.compat.v1 as tf\nfrom tensorflow.contrib import training as contrib_training\n\nflags = tf.app.flags\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string(\n    'master', '',\n    'The TensorFlow master to use.')\nflags.DEFINE_string(\n    'examples_path', None,\n    'Path to a TFRecord file of NoteSequence examples. Overrides the config.')\nflags.DEFINE_string(\n    'tfds_name', None,\n    'TensorFlow Datasets dataset name to use. Overrides the config.')\nflags.DEFINE_string(\n    'run_dir', None,\n    'Path where checkpoints and summary events will be located during '\n    'training and evaluation. Separate subdirectories `train` and `eval` '\n    'will be created within this directory.')\nflags.DEFINE_integer(\n    'num_steps', 200000,\n    'Number of training steps or `None` for infinite.')\nflags.DEFINE_integer(\n    'eval_num_batches', None,\n    'Number of batches to use during evaluation or `None` for all batches '\n    'in the data source.')\nflags.DEFINE_integer(\n    'checkpoints_to_keep', 100,\n    'Maximum number of checkpoints to keep in `train` mode or 0 for infinite.')\nflags.DEFINE_integer(\n    'keep_checkpoint_every_n_hours', 1,\n    'In addition to checkpoints_to_keep, keep a checkpoint every N hours.')\nflags.DEFINE_string(\n    'mode', 'train',\n    'Which mode to use (`train` or `eval`).')\nflags.DEFINE_string(\n    'config', '',\n    'The name of the config to use.')\nflags.DEFINE_string(\n    'hparams', '',\n    'A comma-separated list of `name=value` hyperparameter values to merge '\n    'with those in the config.')\nflags.DEFINE_bool(\n    'cache_dataset', True,\n    'Whether to cache the dataset in memory for improved training speed. May '\n    'cause memory errors for very large datasets.')\nflags.DEFINE_integer(\n    'task', 0,\n    'The task number for this worker.')\nflags.DEFINE_integer(\n    'num_ps_tasks', 0,\n    'The number of parameter server tasks.')\nflags.DEFINE_integer(\n    'num_sync_workers', 0,\n    'The number of synchronized workers.')\nflags.DEFINE_integer(\n    'num_data_threads', 4,\n    'The number of data preprocessing threads.')\nflags.DEFINE_string(\n    'eval_dir_suffix', '',\n    'Suffix to add to eval output directory.')\nflags.DEFINE_string(\n    'log', 'INFO',\n    'The threshold for what messages will be logged: '\n    'DEBUG, INFO, WARN, ERROR, or FATAL.')\n\n\n# Should not be called from within the graph to avoid redundant summaries.\ndef _trial_summary(hparams, examples_path, output_dir):\n  \"\"\"Writes a tensorboard text summary of the trial.\"\"\"\n\n  examples_path_summary = tf.summary.text(\n      'examples_path', tf.constant(examples_path, name='examples_path'),\n      collections=[])\n\n  hparams_dict = hparams.values()\n\n  # Create a markdown table from hparams.\n  header = '| Key | Value |\\n| :--- | :--- |\\n'\n  keys = sorted(hparams_dict.keys())\n  lines = ['| %s | %s |' % (key, str(hparams_dict[key])) for key in keys]\n  hparams_table = header + '\\n'.join(lines) + '\\n'\n\n  hparam_summary = tf.summary.text(\n      'hparams', tf.constant(hparams_table, name='hparams'), collections=[])\n\n  with tf.Session() as sess:\n    writer = tf.summary.FileWriter(output_dir, graph=sess.graph)\n    writer.add_summary(examples_path_summary.eval())\n    writer.add_summary(hparam_summary.eval())\n    writer.close()\n\n\ndef _get_input_tensors(dataset, config):\n  \"\"\"Get input tensors from dataset.\"\"\"\n  batch_size = config.hparams.batch_size\n  iterator = dataset.make_one_shot_iterator()\n  (input_sequence, output_sequence, control_sequence,\n   sequence_length) = iterator.get_next()\n  input_sequence.set_shape(\n      [batch_size, None, config.data_converter.input_depth])\n  output_sequence.set_shape(\n      [batch_size, None, config.data_converter.output_depth])\n  if not config.data_converter.control_depth:\n    control_sequence = None\n  else:\n    control_sequence.set_shape(\n        [batch_size, None, config.data_converter.control_depth])\n  sequence_length.set_shape([batch_size] + sequence_length.shape[1:].as_list())\n\n  return {\n      'input_sequence': input_sequence,\n      'output_sequence': output_sequence,\n      'control_sequence': control_sequence,\n      'sequence_length': sequence_length\n  }\n\n\ndef train(train_dir,\n          config,\n          dataset_fn,\n          checkpoints_to_keep=5,\n          keep_checkpoint_every_n_hours=1,\n          num_steps=None,\n          master='',\n          num_sync_workers=0,\n          num_ps_tasks=0,\n          task=0):\n  \"\"\"Train loop.\"\"\"\n  tf.gfile.MakeDirs(train_dir)\n  is_chief = (task == 0)\n  if is_chief:\n    _trial_summary(\n        config.hparams, config.train_examples_path or config.tfds_name,\n        train_dir)\n  with tf.Graph().as_default():\n    with tf.device(tf.train.replica_device_setter(\n        num_ps_tasks, merge_devices=True)):\n\n      model = config.model\n      model.build(config.hparams,\n                  config.data_converter.output_depth,\n                  is_training=True)\n\n      optimizer = model.train(**_get_input_tensors(dataset_fn(), config))\n\n      hooks = []\n      if num_sync_workers:\n        optimizer = tf.train.SyncReplicasOptimizer(\n            optimizer,\n            num_sync_workers)\n        hooks.append(optimizer.make_session_run_hook(is_chief))\n\n      grads, var_list = list(zip(*optimizer.compute_gradients(model.loss)))\n      global_norm = tf.global_norm(grads)\n      tf.summary.scalar('global_norm', global_norm)\n\n      if config.hparams.clip_mode == 'value':\n        g = config.hparams.grad_clip\n        clipped_grads = [tf.clip_by_value(grad, -g, g) for grad in grads]\n      elif config.hparams.clip_mode == 'global_norm':\n        clipped_grads = tf.cond(\n            global_norm < config.hparams.grad_norm_clip_to_zero,\n            lambda: tf.clip_by_global_norm(  # pylint:disable=g-long-lambda\n                grads, config.hparams.grad_clip, use_norm=global_norm)[0],\n            lambda: [tf.zeros(tf.shape(g)) for g in grads])\n      else:\n        raise ValueError(\n            'Unknown clip_mode: {}'.format(config.hparams.clip_mode))\n      train_op = optimizer.apply_gradients(\n          list(zip(clipped_grads, var_list)),\n          global_step=model.global_step,\n          name='train_step')\n\n      logging_dict = {'global_step': model.global_step,\n                      'loss': model.loss}\n\n      hooks.append(tf.train.LoggingTensorHook(logging_dict, every_n_iter=100))\n      if num_steps:\n        hooks.append(tf.train.StopAtStepHook(last_step=num_steps))\n\n      scaffold = tf.train.Scaffold(\n          saver=tf.train.Saver(\n              max_to_keep=checkpoints_to_keep,\n              keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours))\n      contrib_training.train(\n          train_op=train_op,\n          logdir=train_dir,\n          scaffold=scaffold,\n          hooks=hooks,\n          save_checkpoint_secs=60,\n          master=master,\n          is_chief=is_chief)\n\n\ndef evaluate(train_dir,\n             eval_dir,\n             config,\n             dataset_fn,\n             num_batches,\n             master=''):\n  \"\"\"Evaluate the model repeatedly.\"\"\"\n  tf.gfile.MakeDirs(eval_dir)\n\n  _trial_summary(\n      config.hparams, config.eval_examples_path or config.tfds_name, eval_dir)\n  with tf.Graph().as_default():\n    model = config.model\n    model.build(config.hparams,\n                config.data_converter.output_depth,\n                is_training=False)\n\n    eval_op = model.eval(\n        **_get_input_tensors(dataset_fn().take(num_batches), config))\n\n    hooks = [\n        contrib_training.StopAfterNEvalsHook(num_batches),\n        contrib_training.SummaryAtEndHook(eval_dir)\n    ]\n    contrib_training.evaluate_repeatedly(\n        train_dir,\n        eval_ops=eval_op,\n        hooks=hooks,\n        eval_interval_secs=60,\n        master=master)\n\n\ndef run(config_map,\n        tf_file_reader=tf.data.TFRecordDataset,\n        file_reader=tf.python_io.tf_record_iterator):\n  \"\"\"Load model params, save config file and start trainer.\n\n  Args:\n    config_map: Dictionary mapping configuration name to Config object.\n    tf_file_reader: The tf.data.Dataset class to use for reading files.\n    file_reader: The Python reader to use for reading files.\n\n  Raises:\n    ValueError: if required flags are missing or invalid.\n  \"\"\"\n  if not FLAGS.run_dir:\n    raise ValueError('Invalid run directory: %s' % FLAGS.run_dir)\n  run_dir = os.path.expanduser(FLAGS.run_dir)\n  train_dir = os.path.join(run_dir, 'train')\n\n  if FLAGS.mode not in ['train', 'eval']:\n    raise ValueError('Invalid mode: %s' % FLAGS.mode)\n\n  if FLAGS.config not in config_map:\n    raise ValueError('Invalid config: %s' % FLAGS.config)\n  config = config_map[FLAGS.config]\n  if FLAGS.hparams:\n    config.hparams.parse(FLAGS.hparams)\n  config_update_map = {}\n  if FLAGS.examples_path:\n    config_update_map['%s_examples_path' % FLAGS.mode] = os.path.expanduser(\n        FLAGS.examples_path)\n  if FLAGS.tfds_name:\n    if FLAGS.examples_path:\n      raise ValueError(\n          'At most one of --examples_path and --tfds_name can be set.')\n    config_update_map['tfds_name'] = FLAGS.tfds_name\n    config_update_map['eval_examples_path'] = None\n    config_update_map['train_examples_path'] = None\n  config = configs.update_config(config, config_update_map)\n  if FLAGS.num_sync_workers:\n    config.hparams.batch_size //= FLAGS.num_sync_workers\n\n  if FLAGS.mode == 'train':\n    is_training = True\n  elif FLAGS.mode == 'eval':\n    is_training = False\n  else:\n    raise ValueError('Invalid mode: {}'.format(FLAGS.mode))\n\n  def dataset_fn():\n    return data.get_dataset(\n        config,\n        tf_file_reader=tf_file_reader,\n        num_threads=FLAGS.num_data_threads,\n        is_training=is_training,\n        cache_dataset=FLAGS.cache_dataset)\n\n  if is_training:\n    train(\n        train_dir,\n        config=config,\n        dataset_fn=dataset_fn,\n        checkpoints_to_keep=FLAGS.checkpoints_to_keep,\n        keep_checkpoint_every_n_hours=FLAGS.keep_checkpoint_every_n_hours,\n        num_steps=FLAGS.num_steps,\n        master=FLAGS.master,\n        num_sync_workers=FLAGS.num_sync_workers,\n        num_ps_tasks=FLAGS.num_ps_tasks,\n        task=FLAGS.task)\n  else:\n    num_batches = FLAGS.eval_num_batches or data.count_examples(\n        config.eval_examples_path,\n        config.tfds_name,\n        config.data_converter,\n        file_reader) // config.hparams.batch_size\n    eval_dir = os.path.join(run_dir, 'eval' + FLAGS.eval_dir_suffix)\n    evaluate(\n        train_dir,\n        eval_dir,\n        config=config,\n        dataset_fn=dataset_fn,\n        num_batches=num_batches,\n        master=FLAGS.master)\n\n\ndef main(unused_argv):\n  tf.logging.set_verbosity(FLAGS.log)\n  run(configs.CONFIG_MAP)\n\n\ndef console_entry_point():\n  tf.app.run(main)\n\n\nif __name__ == '__main__':\n  console_entry_point()\n", "path": "magenta/models/music_vae/music_vae_train.py"}], "after_files": [{"content": "# Copyright 2020 The Magenta Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Lint as: python3\n\"\"\"MusicVAE training script.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport os\n\nfrom magenta.models.music_vae import configs\nfrom magenta.models.music_vae import data\nimport tensorflow.compat.v1 as tf\nfrom tensorflow.contrib import training as contrib_training\n\nflags = tf.app.flags\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string(\n    'master', '',\n    'The TensorFlow master to use.')\nflags.DEFINE_string(\n    'examples_path', None,\n    'Path to a TFRecord file of NoteSequence examples. Overrides the config.')\nflags.DEFINE_string(\n    'tfds_name', None,\n    'TensorFlow Datasets dataset name to use. Overrides the config.')\nflags.DEFINE_string(\n    'run_dir', None,\n    'Path where checkpoints and summary events will be located during '\n    'training and evaluation. Separate subdirectories `train` and `eval` '\n    'will be created within this directory.')\nflags.DEFINE_integer(\n    'num_steps', 200000,\n    'Number of training steps or `None` for infinite.')\nflags.DEFINE_integer(\n    'eval_num_batches', None,\n    'Number of batches to use during evaluation or `None` for all batches '\n    'in the data source.')\nflags.DEFINE_integer(\n    'checkpoints_to_keep', 100,\n    'Maximum number of checkpoints to keep in `train` mode or 0 for infinite.')\nflags.DEFINE_integer(\n    'keep_checkpoint_every_n_hours', 1,\n    'In addition to checkpoints_to_keep, keep a checkpoint every N hours.')\nflags.DEFINE_string(\n    'mode', 'train',\n    'Which mode to use (`train` or `eval`).')\nflags.DEFINE_string(\n    'config', '',\n    'The name of the config to use.')\nflags.DEFINE_string(\n    'hparams', '',\n    'A comma-separated list of `name=value` hyperparameter values to merge '\n    'with those in the config.')\nflags.DEFINE_bool(\n    'cache_dataset', True,\n    'Whether to cache the dataset in memory for improved training speed. May '\n    'cause memory errors for very large datasets.')\nflags.DEFINE_integer(\n    'task', 0,\n    'The task number for this worker.')\nflags.DEFINE_integer(\n    'num_ps_tasks', 0,\n    'The number of parameter server tasks.')\nflags.DEFINE_integer(\n    'num_sync_workers', 0,\n    'The number of synchronized workers.')\nflags.DEFINE_integer(\n    'num_data_threads', 4,\n    'The number of data preprocessing threads.')\nflags.DEFINE_string(\n    'eval_dir_suffix', '',\n    'Suffix to add to eval output directory.')\nflags.DEFINE_string(\n    'log', 'INFO',\n    'The threshold for what messages will be logged: '\n    'DEBUG, INFO, WARN, ERROR, or FATAL.')\n\n\n# Should not be called from within the graph to avoid redundant summaries.\ndef _trial_summary(hparams, examples_path, output_dir):\n  \"\"\"Writes a tensorboard text summary of the trial.\"\"\"\n\n  examples_path_summary = tf.summary.text(\n      'examples_path', tf.constant(examples_path, name='examples_path'),\n      collections=[])\n\n  hparams_dict = hparams.values()\n\n  # Create a markdown table from hparams.\n  header = '| Key | Value |\\n| :--- | :--- |\\n'\n  keys = sorted(hparams_dict.keys())\n  lines = ['| %s | %s |' % (key, str(hparams_dict[key])) for key in keys]\n  hparams_table = header + '\\n'.join(lines) + '\\n'\n\n  hparam_summary = tf.summary.text(\n      'hparams', tf.constant(hparams_table, name='hparams'), collections=[])\n\n  with tf.Session() as sess:\n    writer = tf.summary.FileWriter(output_dir, graph=sess.graph)\n    writer.add_summary(examples_path_summary.eval())\n    writer.add_summary(hparam_summary.eval())\n    writer.close()\n\n\ndef _get_input_tensors(dataset, config):\n  \"\"\"Get input tensors from dataset.\"\"\"\n  batch_size = config.hparams.batch_size\n  iterator = tf.data.make_one_shot_iterator(dataset)\n  (input_sequence, output_sequence, control_sequence,\n   sequence_length) = iterator.get_next()\n  input_sequence.set_shape(\n      [batch_size, None, config.data_converter.input_depth])\n  output_sequence.set_shape(\n      [batch_size, None, config.data_converter.output_depth])\n  if not config.data_converter.control_depth:\n    control_sequence = None\n  else:\n    control_sequence.set_shape(\n        [batch_size, None, config.data_converter.control_depth])\n  sequence_length.set_shape([batch_size] + sequence_length.shape[1:].as_list())\n\n  return {\n      'input_sequence': input_sequence,\n      'output_sequence': output_sequence,\n      'control_sequence': control_sequence,\n      'sequence_length': sequence_length\n  }\n\n\ndef train(train_dir,\n          config,\n          dataset_fn,\n          checkpoints_to_keep=5,\n          keep_checkpoint_every_n_hours=1,\n          num_steps=None,\n          master='',\n          num_sync_workers=0,\n          num_ps_tasks=0,\n          task=0):\n  \"\"\"Train loop.\"\"\"\n  tf.gfile.MakeDirs(train_dir)\n  is_chief = (task == 0)\n  if is_chief:\n    _trial_summary(\n        config.hparams, config.train_examples_path or config.tfds_name,\n        train_dir)\n  with tf.Graph().as_default():\n    with tf.device(tf.train.replica_device_setter(\n        num_ps_tasks, merge_devices=True)):\n\n      model = config.model\n      model.build(config.hparams,\n                  config.data_converter.output_depth,\n                  is_training=True)\n\n      optimizer = model.train(**_get_input_tensors(dataset_fn(), config))\n\n      hooks = []\n      if num_sync_workers:\n        optimizer = tf.train.SyncReplicasOptimizer(\n            optimizer,\n            num_sync_workers)\n        hooks.append(optimizer.make_session_run_hook(is_chief))\n\n      grads, var_list = list(zip(*optimizer.compute_gradients(model.loss)))\n      global_norm = tf.global_norm(grads)\n      tf.summary.scalar('global_norm', global_norm)\n\n      if config.hparams.clip_mode == 'value':\n        g = config.hparams.grad_clip\n        clipped_grads = [tf.clip_by_value(grad, -g, g) for grad in grads]\n      elif config.hparams.clip_mode == 'global_norm':\n        clipped_grads = tf.cond(\n            global_norm < config.hparams.grad_norm_clip_to_zero,\n            lambda: tf.clip_by_global_norm(  # pylint:disable=g-long-lambda\n                grads, config.hparams.grad_clip, use_norm=global_norm)[0],\n            lambda: [tf.zeros(tf.shape(g)) for g in grads])\n      else:\n        raise ValueError(\n            'Unknown clip_mode: {}'.format(config.hparams.clip_mode))\n      train_op = optimizer.apply_gradients(\n          list(zip(clipped_grads, var_list)),\n          global_step=model.global_step,\n          name='train_step')\n\n      logging_dict = {'global_step': model.global_step,\n                      'loss': model.loss}\n\n      hooks.append(tf.train.LoggingTensorHook(logging_dict, every_n_iter=100))\n      if num_steps:\n        hooks.append(tf.train.StopAtStepHook(last_step=num_steps))\n\n      scaffold = tf.train.Scaffold(\n          saver=tf.train.Saver(\n              max_to_keep=checkpoints_to_keep,\n              keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours))\n      contrib_training.train(\n          train_op=train_op,\n          logdir=train_dir,\n          scaffold=scaffold,\n          hooks=hooks,\n          save_checkpoint_secs=60,\n          master=master,\n          is_chief=is_chief)\n\n\ndef evaluate(train_dir,\n             eval_dir,\n             config,\n             dataset_fn,\n             num_batches,\n             master=''):\n  \"\"\"Evaluate the model repeatedly.\"\"\"\n  tf.gfile.MakeDirs(eval_dir)\n\n  _trial_summary(\n      config.hparams, config.eval_examples_path or config.tfds_name, eval_dir)\n  with tf.Graph().as_default():\n    model = config.model\n    model.build(config.hparams,\n                config.data_converter.output_depth,\n                is_training=False)\n\n    eval_op = model.eval(\n        **_get_input_tensors(dataset_fn().take(num_batches), config))\n\n    hooks = [\n        contrib_training.StopAfterNEvalsHook(num_batches),\n        contrib_training.SummaryAtEndHook(eval_dir)\n    ]\n    contrib_training.evaluate_repeatedly(\n        train_dir,\n        eval_ops=eval_op,\n        hooks=hooks,\n        eval_interval_secs=60,\n        master=master)\n\n\ndef run(config_map,\n        tf_file_reader=tf.data.TFRecordDataset,\n        file_reader=tf.python_io.tf_record_iterator):\n  \"\"\"Load model params, save config file and start trainer.\n\n  Args:\n    config_map: Dictionary mapping configuration name to Config object.\n    tf_file_reader: The tf.data.Dataset class to use for reading files.\n    file_reader: The Python reader to use for reading files.\n\n  Raises:\n    ValueError: if required flags are missing or invalid.\n  \"\"\"\n  if not FLAGS.run_dir:\n    raise ValueError('Invalid run directory: %s' % FLAGS.run_dir)\n  run_dir = os.path.expanduser(FLAGS.run_dir)\n  train_dir = os.path.join(run_dir, 'train')\n\n  if FLAGS.mode not in ['train', 'eval']:\n    raise ValueError('Invalid mode: %s' % FLAGS.mode)\n\n  if FLAGS.config not in config_map:\n    raise ValueError('Invalid config: %s' % FLAGS.config)\n  config = config_map[FLAGS.config]\n  if FLAGS.hparams:\n    config.hparams.parse(FLAGS.hparams)\n  config_update_map = {}\n  if FLAGS.examples_path:\n    config_update_map['%s_examples_path' % FLAGS.mode] = os.path.expanduser(\n        FLAGS.examples_path)\n  if FLAGS.tfds_name:\n    if FLAGS.examples_path:\n      raise ValueError(\n          'At most one of --examples_path and --tfds_name can be set.')\n    config_update_map['tfds_name'] = FLAGS.tfds_name\n    config_update_map['eval_examples_path'] = None\n    config_update_map['train_examples_path'] = None\n  config = configs.update_config(config, config_update_map)\n  if FLAGS.num_sync_workers:\n    config.hparams.batch_size //= FLAGS.num_sync_workers\n\n  if FLAGS.mode == 'train':\n    is_training = True\n  elif FLAGS.mode == 'eval':\n    is_training = False\n  else:\n    raise ValueError('Invalid mode: {}'.format(FLAGS.mode))\n\n  def dataset_fn():\n    return data.get_dataset(\n        config,\n        tf_file_reader=tf_file_reader,\n        num_threads=FLAGS.num_data_threads,\n        is_training=is_training,\n        cache_dataset=FLAGS.cache_dataset)\n\n  if is_training:\n    train(\n        train_dir,\n        config=config,\n        dataset_fn=dataset_fn,\n        checkpoints_to_keep=FLAGS.checkpoints_to_keep,\n        keep_checkpoint_every_n_hours=FLAGS.keep_checkpoint_every_n_hours,\n        num_steps=FLAGS.num_steps,\n        master=FLAGS.master,\n        num_sync_workers=FLAGS.num_sync_workers,\n        num_ps_tasks=FLAGS.num_ps_tasks,\n        task=FLAGS.task)\n  else:\n    num_batches = FLAGS.eval_num_batches or data.count_examples(\n        config.eval_examples_path,\n        config.tfds_name,\n        config.data_converter,\n        file_reader) // config.hparams.batch_size\n    eval_dir = os.path.join(run_dir, 'eval' + FLAGS.eval_dir_suffix)\n    evaluate(\n        train_dir,\n        eval_dir,\n        config=config,\n        dataset_fn=dataset_fn,\n        num_batches=num_batches,\n        master=FLAGS.master)\n\n\ndef main(unused_argv):\n  tf.logging.set_verbosity(FLAGS.log)\n  run(configs.CONFIG_MAP)\n\n\ndef console_entry_point():\n  tf.app.run(main)\n\n\nif __name__ == '__main__':\n  console_entry_point()\n", "path": "magenta/models/music_vae/music_vae_train.py"}]}