Merge pull request #5546 from cshallue/master

Improvements to AstroNet and add AstroWaveNet

research/astronet/astronet/ops/training.py

@@ -51,7 +51,7 @@ def create_optimizer(hparams, learning_rate, use_tpu=False):
     hparams: ConfigDict containing the optimizer configuration.
     learning_rate: A Python float or a scalar Tensor.
     use_tpu: If True, the returned optimizer is wrapped in a
-        CrossShardOptimizer.
+      CrossShardOptimizer.
 
   Returns:
     A TensorFlow optimizer.
@@ -74,7 +74,7 @@ def create_optimizer(hparams, learning_rate, use_tpu=False):
   elif optimizer_name == "rmsprop":
     optimizer = tf.RMSPropOptimizer(learning_rate)
   else:
-    raise ValueError("Unknown optimizer: %s" % hparams.optimizer)
+    raise ValueError("Unknown optimizer: {}".format(hparams.optimizer))
 
   if use_tpu:
     optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)

research/astronet/astronet/train.py

@@ -26,6 +26,7 @@ import tensorflow as tf
 from astronet import models
 from astronet.util import config_util
 from astronet.util import configdict
+from astronet.util import estimator_runner
 from astronet.util import estimator_util
 
 parser = argparse.ArgumentParser()
@@ -112,11 +113,14 @@ def main(_):
         file_pattern=FLAGS.eval_files,
         input_config=config.inputs,
         mode=tf.estimator.ModeKeys.EVAL)
+    eval_args = {
+        "val": (eval_input_fn, None)  # eval_name: (input_fn, eval_steps)
+    }
 
-    for _ in estimator_util.continuous_train_and_eval(
+    for _ in estimator_runner.continuous_train_and_eval(
         estimator=estimator,
         train_input_fn=train_input_fn,
-        eval_input_fn=eval_input_fn,
+        eval_args=eval_args,
         train_steps=FLAGS.train_steps):
       # continuous_train_and_eval() yields evaluation metrics after each
       # training epoch. We don't do anything here.

research/astronet/astronet/util/BUILD

@@ -32,6 +32,12 @@ py_test(
     deps = [":config_util"],
 )
 
+py_library(
+    name = "estimator_runner",
+    srcs = ["estimator_runner.py"],
+    srcs_version = "PY2AND3",
+)
+
 py_library(
     name = "estimator_util",
     srcs = ["estimator_util.py"],
@@ -47,6 +53,7 @@ py_library(
     name = "example_util",
     srcs = ["example_util.py"],
     srcs_version = "PY2AND3",
+    visibility = ["//visibility:public"],
 )
 
 py_test(

research/astronet/astronet/util/config_util.py

@@ -49,19 +49,28 @@ def parse_json(json_string_or_file):
       with tf.gfile.Open(json_string_or_file) as f:
         json_dict = json.load(f)
     except ValueError as json_file_parsing_error:
-      raise ValueError("Unable to parse the content of the json file %s. "
-                       "Parsing error: %s." % (json_string_or_file,
-                                               json_file_parsing_error.message))
+      raise ValueError("Unable to parse the content of the json file {}. "
+                       "Parsing error: {}.".format(
+                           json_string_or_file,
+                           json_file_parsing_error.message))
     except tf.gfile.FileError:
       message = ("Unable to parse the input parameter neither as literal "
                  "JSON nor as the name of a file that exists.\n"
-                 "JSON parsing error: %s\n\n Input parameter:\n%s." %
-                 (literal_json_parsing_error.message, json_string_or_file))
+                 "JSON parsing error: {}\n\n Input parameter:\n{}.".format(
+                     literal_json_parsing_error.message, json_string_or_file))
       raise ValueError(message)
 
   return json_dict
 
 
+def to_json(config):
+  """Converts a JSON-serializable configuration object to a JSON string."""
+  if hasattr(config, "to_json") and callable(config.to_json):
+    return config.to_json(indent=2)
+  else:
+    return json.dumps(config, indent=2)
+
+
 def log_and_save_config(config, output_dir):
   """Logs and writes a JSON-serializable configuration object.
 
@@ -69,10 +78,7 @@ def log_and_save_config(config, output_dir):
     config: A JSON-serializable object.
     output_dir: Destination directory.
   """
-  if hasattr(config, "to_json") and callable(config.to_json):
-    config_json = config.to_json(indent=2)
-  else:
-    config_json = json.dumps(config, indent=2)
+  config_json = to_json(config)
   tf.logging.info("config: %s", config_json)
 
   tf.gfile.MakeDirs(output_dir)
@@ -104,7 +110,7 @@ def unflatten(flat_config):
 
   Args:
     flat_config: A dictionary with strings as keys where nested configuration
-        parameters are represented with period-separated names.
+      parameters are represented with period-separated names.
 
   Returns:
     A dictionary nested according to the keys of the input dictionary.

research/astronet/astronet/util/estimator_runner.py

+# Copyright 2018 The TensorFlow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Functions for training and evaluation using a TensorFlow Estimator."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+def evaluate(estimator, eval_args):
+  """Runs evaluation on the latest model checkpoint.
+
+  Args:
+    estimator: Instance of tf.Estimator.
+    eval_args: Dictionary of {eval_name: (input_fn, eval_steps)} where eval_name
+      is the name of the evaluation set (e.g. "train" or "val"), input_fn is an
+      input function returning a tuple (features, labels), and eval_steps is the
+      number of steps for which to evaluate the model (if None, evaluates until
+      input_fn raises an end-of-input exception).
+
+  Returns:
+    global_step: The global step of the checkpoint evaluated.
+    values: A dict of metric values from the evaluation. May be empty, e.g. if
+        the training job has not yet saved a checkpoint or the checkpoint is
+        deleted by the time the TPU worker initializes.
+  """
+  # Default return values if evaluation fails.
+  global_step = None
+  values = {}
+
+  latest_checkpoint = estimator.latest_checkpoint()
+  if not latest_checkpoint:
+    # This is expected if the training job has not yet saved a checkpoint.
+    return global_step, values
+
+  tf.logging.info("Starting evaluation on checkpoint %s", latest_checkpoint)
+  try:
+    for eval_name, (input_fn, eval_steps) in eval_args.items():
+      values[eval_name] = estimator.evaluate(
+          input_fn, steps=eval_steps, name=eval_name)
+      if global_step is None:
+        global_step = values[eval_name].get("global_step")
+  except (tf.errors.NotFoundError, ValueError):
+    # Expected under some conditions, e.g. checkpoint is already deleted by the
+    # trainer process. Increasing RunConfig.keep_checkpoint_max may prevent this
+    # in some cases.
+    tf.logging.info("Checkpoint %s no longer exists, skipping evaluation.",
+                    latest_checkpoint)
+
+  return global_step, values
+
+
+def continuous_eval(estimator,
+                    eval_args,
+                    train_steps=None,
+                    timeout_secs=None,
+                    timeout_fn=None):
+  """Runs evaluation whenever there's a new checkpoint.
+
+  Args:
+    estimator: Instance of tf.Estimator.
+    eval_args: Dictionary of {eval_name: (input_fn, eval_steps)} where eval_name
+      is the name of the evaluation set (e.g. "train" or "val"), input_fn is an
+      input function returning a tuple (features, labels), and eval_steps is the
+      number of steps for which to evaluate the model (if None, evaluates until
+      input_fn raises an end-of-input exception).
+    train_steps: The number of steps the model will train for. This function
+      will terminate once the model has finished training.
+    timeout_secs: Number of seconds to wait for new checkpoints. If None, wait
+      indefinitely.
+    timeout_fn: Optional function to call after timeout. The iterator will exit
+      if and only if the function returns True.
+
+  Yields:
+    A dict of metric values from each evaluation. May be empty, e.g. if the
+    training job has not yet saved a checkpoint or the checkpoint is deleted by
+    the time the TPU worker initializes.
+  """
+  for _ in tf.contrib.training.checkpoints_iterator(
+      estimator.model_dir, timeout=timeout_secs, timeout_fn=timeout_fn):
+    global_step, values = evaluate(estimator, eval_args)
+    yield global_step, values
+
+    global_step = global_step or 0  # Ensure global_step is not None.
+    if train_steps and global_step >= train_steps:
+      break
+
+
+def continuous_train_and_eval(estimator,
+                              train_input_fn,
+                              eval_args,
+                              local_eval_frequency=None,
+                              train_hooks=None,
+                              train_steps=None):
+  """Alternates training and evaluation.
+
+  Args:
+    estimator: Instance of tf.Estimator.
+    train_input_fn: Input function returning a tuple (features, labels).
+    eval_args: Dictionary of {eval_name: (input_fn, eval_steps)} where eval_name
+      is the name of the evaluation set (e.g. "train" or "val"), input_fn is an
+      input function returning a tuple (features, labels), and eval_steps is the
+      number of steps for which to evaluate the model (if None, evaluates until
+      input_fn raises an end-of-input exception).
+    local_eval_frequency: The number of training steps between evaluations. If
+      None, trains until train_input_fn raises an end-of-input exception.
+    train_hooks: List of SessionRunHook subclass instances. Used for callbacks
+      inside the training call.
+    train_steps: The total number of steps to train the model for.
+
+  Yields:
+    A dict of metric values from each evaluation. May be empty, e.g. if the
+    training job has not yet saved a checkpoint or the checkpoint is deleted by
+    the time the TPU worker initializes.
+  """
+  while True:
+    # We run evaluation before training in this loop to prevent evaluation from
+    # being skipped if the process is interrupted.
+    global_step, values = evaluate(estimator, eval_args)
+    yield global_step, values
+
+    global_step = global_step or 0  # Ensure global_step is not None.
+    if train_steps and global_step >= train_steps:
+      break
+
+    # Decide how many steps before the next evaluation.
+    steps = local_eval_frequency
+    if train_steps:
+      remaining_steps = train_steps - global_step
+      steps = min(steps, remaining_steps) if steps else remaining_steps
+
+    tf.logging.info("Starting training at global step %d", global_step)
+    estimator.train(train_input_fn, hooks=train_hooks, steps=steps)

research/astronet/astronet/util/estimator_util.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Functions for training models with the TensorFlow Estimator API."""
+"""Helper functions for creating a TensorFlow Estimator."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -27,71 +27,104 @@ from astronet.ops import metrics
 from astronet.ops import training
 
 
-def create_input_fn(file_pattern,
-                    input_config,
-                    mode,
-                    shuffle_values_buffer=0,
-                    repeat=1):
-  """Creates an input_fn that reads a dataset from sharded TFRecord files.
+class _InputFn(object):
+  """Class that acts as a callable input function for Estimator train / eval."""
 
-  Args:
-    file_pattern: File pattern matching input TFRecord files, e.g.
+  def __init__(self,
+               file_pattern,
+               input_config,
+               mode,
+               shuffle_values_buffer=0,
+               repeat=1):
+    """Initializes the input function.
+
+    Args:
+      file_pattern: File pattern matching input TFRecord files, e.g.
         "/tmp/train-?????-of-00100". May also be a comma-separated list of file
         patterns.
-    input_config: ConfigDict containing feature and label specifications.
-    mode: A tf.estimator.ModeKeys.
-    shuffle_values_buffer: If > 0, shuffle examples using a buffer of this size.
-    repeat: The number of times to repeat the dataset. If None or -1 the
+      input_config: ConfigDict containing feature and label specifications.
+      mode: A tf.estimator.ModeKeys.
+      shuffle_values_buffer: If > 0, shuffle examples using a buffer of this
+        size.
+      repeat: The number of times to repeat the dataset. If None or -1 the
         elements will be repeated indefinitely.
-
-  Returns:
-    A callable that builds an input pipeline and returns (features, labels).
-  """
-  include_labels = (
-      mode in [tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL])
-  reverse_time_series_prob = 0.5 if mode == tf.estimator.ModeKeys.TRAIN else 0
-  shuffle_filenames = (mode == tf.estimator.ModeKeys.TRAIN)
-
-  def input_fn(config, params):
-    """Builds an input pipeline that reads a dataset from TFRecord files."""
+    """
+    self._file_pattern = file_pattern
+    self._input_config = input_config
+    self._mode = mode
+    self._shuffle_values_buffer = shuffle_values_buffer
+    self._repeat = repeat
+
+  def __call__(self, config, params):
+    """Builds the input pipeline."""
     # Infer whether this input_fn was called by Estimator or TPUEstimator using
     # the config type.
     use_tpu = isinstance(config, tf.contrib.tpu.RunConfig)
 
+    mode = self._mode
+    include_labels = (
+        mode in [tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL])
+    reverse_time_series_prob = 0.5 if mode == tf.estimator.ModeKeys.TRAIN else 0
+    shuffle_filenames = (mode == tf.estimator.ModeKeys.TRAIN)
     dataset = dataset_ops.build_dataset(
-        file_pattern=file_pattern,
-        input_config=input_config,
+        file_pattern=self._file_pattern,
+        input_config=self._input_config,
         batch_size=params["batch_size"],
         include_labels=include_labels,
         reverse_time_series_prob=reverse_time_series_prob,
         shuffle_filenames=shuffle_filenames,
-        shuffle_values_buffer=shuffle_values_buffer,
-        repeat=repeat,
+        shuffle_values_buffer=self._shuffle_values_buffer,
+        repeat=self._repeat,
         use_tpu=use_tpu)
 
     return dataset
 
-  return input_fn
 
-
-def create_model_fn(model_class, hparams, use_tpu=False):
-  """Wraps model_class as an Estimator or TPUEstimator model_fn.
+def create_input_fn(file_pattern,
+                    input_config,
+                    mode,
+                    shuffle_values_buffer=0,
+                    repeat=1):
+  """Creates an input_fn that reads a dataset from sharded TFRecord files.
 
   Args:
-    model_class: AstroModel or a subclass.
-    hparams: ConfigDict of configuration parameters for building the model.
-    use_tpu: If True, a TPUEstimator model_fn is returned. Otherwise an
-        Estimator model_fn is returned.
+    file_pattern: File pattern matching input TFRecord files, e.g.
+      "/tmp/train-?????-of-00100". May also be a comma-separated list of file
+      patterns.
+    input_config: ConfigDict containing feature and label specifications.
+    mode: A tf.estimator.ModeKeys.
+    shuffle_values_buffer: If > 0, shuffle examples using a buffer of this size.
+    repeat: The number of times to repeat the dataset. If None or -1 the
+      elements will be repeated indefinitely.
 
   Returns:
-    model_fn: A callable that constructs the model and returns a
-        TPUEstimatorSpec if use_tpu is True, otherwise an EstimatorSpec.
+    A callable that builds the input pipeline and returns a tf.data.Dataset
+    object.
   """
-  hparams = copy.deepcopy(hparams)
+  return _InputFn(file_pattern, input_config, mode, shuffle_values_buffer,
+                  repeat)
+
 
-  def model_fn(features, labels, mode, params):
+class _ModelFn(object):
+  """Class that acts as a callable model function for Estimator train / eval."""
+
+  def __init__(self, model_class, hparams, use_tpu=False):
+    """Initializes the model function.
+
+    Args:
+      model_class: Model class.
+      hparams: ConfigDict containing hyperparameters for building and training
+        the model.
+      use_tpu: If True, a TPUEstimator will be returned. Otherwise an Estimator
+        will be returned.
+    """
+    self._model_class = model_class
+    self._base_hparams = hparams
+    self._use_tpu = use_tpu
+
+  def __call__(self, features, labels, mode, params):
     """Builds the model and returns an EstimatorSpec or TPUEstimatorSpec."""
-    # For TPUEstimator, params contains the batch size per TPU core.
+    hparams = copy.deepcopy(self._base_hparams)
     if "batch_size" in params:
       hparams.batch_size = params["batch_size"]
 
@@ -99,14 +132,15 @@ def create_model_fn(model_class, hparams, use_tpu=False):
     if "labels" in features:
       if labels is not None and labels is not features["labels"]:
         raise ValueError(
-            "Conflicting labels: features['labels'] = %s, labels = %s" %
-            (features["labels"], labels))
+            "Conflicting labels: features['labels'] = {}, labels = {}".format(
+                features["labels"], labels))
       labels = features.pop("labels")
 
-    model = model_class(features, labels, hparams, mode)
+    model = self._model_class(features, labels, hparams, mode)
     model.build()
 
     # Possibly create train_op.
+    use_tpu = self._use_tpu
     train_op = None
     if mode == tf.estimator.ModeKeys.TRAIN:
       learning_rate = training.create_learning_rate(hparams, model.global_step)
@@ -137,7 +171,21 @@ def create_model_fn(model_class, hparams, use_tpu=False):
 
     return estimator
 
-  return model_fn
+
+def create_model_fn(model_class, hparams, use_tpu=False):
+  """Wraps model_class as an Estimator or TPUEstimator model_fn.
+
+  Args:
+    model_class: AstroModel or a subclass.
+    hparams: ConfigDict of configuration parameters for building the model.
+    use_tpu: If True, a TPUEstimator model_fn is returned. Otherwise an
+      Estimator model_fn is returned.
+
+  Returns:
+    model_fn: A callable that constructs the model and returns a
+        TPUEstimatorSpec if use_tpu is True, otherwise an EstimatorSpec.
+  """
+  return _ModelFn(model_class, hparams, use_tpu)
 
 
 def create_estimator(model_class,
@@ -155,10 +203,10 @@ def create_estimator(model_class,
     hparams: ConfigDict of configuration parameters for building the model.
     run_config: Optional tf.estimator.RunConfig or tf.contrib.tpu.RunConfig.
     model_dir: Optional directory for saving the model. If not passed
-        explicitly, it must be specified in run_config.
+      explicitly, it must be specified in run_config.
     eval_batch_size: Optional batch size for evaluation on TPU. Only applicable
-        if run_config is a tf.contrib.tpu.RunConfig. Defaults to
-        hparams.batch_size.
+      if run_config is a tf.contrib.tpu.RunConfig. Defaults to
+      hparams.batch_size.
 
   Returns:
     An Estimator object if run_config is None or a tf.estimator.RunConfig, or a
@@ -202,117 +250,3 @@ def create_estimator(model_class,
         params={"batch_size": hparams.batch_size})
 
   return estimator
-
-
-def evaluate(estimator, input_fn, eval_steps=None, eval_name="val"):
-  """Runs evaluation on the latest model checkpoint.
-
-  Args:
-    estimator: Instance of tf.Estimator.
-    input_fn: Input function returning a tuple (features, labels).
-    eval_steps: The number of steps for which to evaluate the model. If None,
-        evaluates until input_fn raises an end-of-input exception.
-    eval_name: Name of the evaluation set, e.g. "train" or "val".
-
-  Returns:
-    A dict of metric values from the evaluation. May be empty, e.g. if the
-    training job has not yet saved a checkpoint or the checkpoint is deleted by
-    the time the TPU worker initializes.
-  """
-  values = {}  # Default return value if evaluation fails.
-
-  latest_checkpoint = tf.train.latest_checkpoint(estimator.model_dir)
-  if not latest_checkpoint:
-    # This is expected if the training job has not yet saved a checkpoint.
-    return values
-
-  tf.logging.info("Starting evaluation on checkpoint %s", latest_checkpoint)
-  try:
-    values = estimator.evaluate(input_fn, steps=eval_steps, name=eval_name)
-  except tf.errors.NotFoundError:
-    # Expected under some conditions, e.g. TPU worker does not finish
-    # initializing until long after the CPU job tells it to start evaluating
-    # and the checkpoint file is deleted already.
-    tf.logging.info("Checkpoint %s no longer exists, skipping evaluation",
-                    latest_checkpoint)
-
-  return values
-
-
-def continuous_eval(estimator,
-                    input_fn,
-                    train_steps=None,
-                    eval_steps=None,
-                    eval_name="val"):
-  """Runs evaluation whenever there's a new checkpoint.
-
-  Args:
-    estimator: Instance of tf.Estimator.
-    input_fn: Input function returning a tuple (features, labels).
-    train_steps: The number of steps the model will train for. This function
-        will terminate once the model has finished training. If None, this
-        function will run forever.
-    eval_steps: The number of steps for which to evaluate the model. If None,
-        evaluates until input_fn raises an end-of-input exception.
-    eval_name: Name of the evaluation set, e.g. "train" or "val".
-
-  Yields:
-    A dict of metric values from each evaluation. May be empty, e.g. if the
-    training job has not yet saved a checkpoint or the checkpoint is deleted by
-    the time the TPU worker initializes.
-  """
-  for _ in tf.contrib.training.checkpoints_iterator(estimator.model_dir):
-    values = evaluate(estimator, input_fn, eval_steps, eval_name)
-    yield values
-
-    global_step = values.get("global_step", 0)
-    if train_steps and global_step >= train_steps:
-      break
-
-
-def continuous_train_and_eval(estimator,
-                              train_input_fn,
-                              eval_input_fn,
-                              local_eval_frequency=None,
-                              train_hooks=None,
-                              train_steps=None,
-                              eval_steps=None,
-                              eval_name="val"):
-  """Alternates training and evaluation.
-
-  Args:
-    estimator: Instance of tf.Estimator.
-    train_input_fn: Input function returning a tuple (features, labels).
-    eval_input_fn: Input function returning a tuple (features, labels).
-    local_eval_frequency: The number of training steps between evaluations. If
-        None, trains until train_input_fn raises an end-of-input exception.
-    train_hooks: List of SessionRunHook subclass instances. Used for callbacks
-        inside the training call.
-    train_steps: The total number of steps to train the model for.
-    eval_steps: The number of steps for which to evaluate the model. If None,
-        evaluates until eval_input_fn raises an end-of-input exception.
-    eval_name: Name of the evaluation set, e.g. "train" or "val".
-
-  Yields:
-    A dict of metric values from each evaluation. May be empty, e.g. if the
-    training job has not yet saved a checkpoint or the checkpoint is deleted by
-    the time the TPU worker initializes.
-  """
-  while True:
-    # We run evaluation before training in this loop to prevent evaluation from
-    # being skipped if the process is interrupted.
-    values = evaluate(estimator, eval_input_fn, eval_steps, eval_name)
-    yield values
-
-    global_step = values.get("global_step", 0)
-    if train_steps and global_step >= train_steps:
-      break
-
-    # Decide how many steps before the next evaluation.
-    steps = local_eval_frequency
-    if train_steps:
-      remaining_steps = train_steps - global_step
-      steps = min(steps, remaining_steps) if steps else remaining_steps
-
-    tf.logging.info("Starting training at global step %d", global_step)
-    estimator.train(train_input_fn, hooks=train_hooks, steps=steps)

research/astronet/astronet/util/example_util.py

@@ -28,7 +28,7 @@ def get_feature(ex, name, kind=None, strict=True):
     ex: A tf.train.Example.
     name: Name of the feature to look up.
     kind: Optional: one of 'bytes_list', 'float_list', 'int64_list'. Inferred if
-        not specified.
+      not specified.
     strict: Whether to raise a KeyError if there is no such feature.
 
   Returns:
@@ -48,7 +48,8 @@ def get_feature(ex, name, kind=None, strict=True):
     return np.array([])  # Feature exists, but it's empty.
 
   if kind and kind != inferred_kind:
-    raise TypeError("Requested %s, but Feature has %s" % (kind, inferred_kind))
+    raise TypeError("Requested {}, but Feature has {}".format(
+        kind, inferred_kind))
 
   return np.array(getattr(ex.features.feature[name], inferred_kind).value)
 
@@ -79,7 +80,12 @@ def _infer_kind(value):
     return "bytes_list"
 
 
-def set_feature(ex, name, value, kind=None, allow_overwrite=False):
+def set_feature(ex,
+                name,
+                value,
+                kind=None,
+                allow_overwrite=False,
+                bytes_encoding="latin-1"):
   """Sets a feature value in a tf.train.Example.
 
   Args:
@@ -87,8 +93,9 @@ def set_feature(ex, name, value, kind=None, allow_overwrite=False):
     name: Name of the feature to set.
     value: Feature value to set. Must be a sequence.
     kind: Optional: one of 'bytes_list', 'float_list', 'int64_list'. Inferred if
-        not specified.
+      not specified.
     allow_overwrite: Whether to overwrite the existing value of the feature.
+    bytes_encoding: Codec for encoding strings when kind = 'bytes_list'.
 
   Raises:
     ValueError: If `allow_overwrite` is False and the feature already exists, or
@@ -99,19 +106,20 @@ def set_feature(ex, name, value, kind=None, allow_overwrite=False):
       del ex.features.feature[name]
     else:
       raise ValueError(
-          "Attempting to set duplicate feature with name: %s" % name)
+          "Attempting to overwrite feature with name: {}. "
+          "Set allow_overwrite=True if this is desired.".format(name))
 
   if not kind:
     kind = _infer_kind(value)
 
   if kind == "bytes_list":
-    value = [str(v).encode("latin-1") for v in value]
+    value = [str(v).encode(bytes_encoding) for v in value]
   elif kind == "float_list":
     value = [float(v) for v in value]
   elif kind == "int64_list":
     value = [int(v) for v in value]
   else:
-    raise ValueError("Unrecognized kind: %s" % kind)
+    raise ValueError("Unrecognized kind: {}".format(kind))
 
   getattr(ex.features.feature[name], kind).value.extend(value)
 
@@ -121,9 +129,13 @@ def set_float_feature(ex, name, value, allow_overwrite=False):
   set_feature(ex, name, value, "float_list", allow_overwrite)
 
 
-def set_bytes_feature(ex, name, value, allow_overwrite=False):
+def set_bytes_feature(ex,
+                      name,
+                      value,
+                      allow_overwrite=False,
+                      bytes_encoding="latin-1"):
   """Sets the value of a bytes feature in a tf.train.Example."""
-  set_feature(ex, name, value, "bytes_list", allow_overwrite)
+  set_feature(ex, name, value, "bytes_list", allow_overwrite, bytes_encoding)
 
 
 def set_int64_feature(ex, name, value, allow_overwrite=False):

research/astronet/astrowavenet/BUILD

+"""A TensorFlow model for generative modeling of light curves."""
+
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])  # Apache 2.0
+
+py_binary(
+    name = "trainer",
+    srcs = ["trainer.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        ":astrowavenet_model",
+        ":configurations",
+        "//astronet/util:config_util",
+        "//astronet/util:configdict",
+        "//astronet/util:estimator_runner",
+        "//astrowavenet/data:kepler_light_curves",
+        "//astrowavenet/data:synthetic_transits",
+        "//astrowavenet/util:estimator_util",
+    ],
+)
+
+py_library(
+    name = "configurations",
+    srcs = ["configurations.py"],
+    srcs_version = "PY2AND3",
+)
+
+py_library(
+    name = "astrowavenet_model",
+    srcs = [
+        "astrowavenet_model.py",
+    ],
+    srcs_version = "PY2AND3",
+)
+
+py_test(
+    name = "astrowavenet_model_test",
+    size = "small",
+    srcs = [
+        "astrowavenet_model_test.py",
+    ],
+    srcs_version = "PY2AND3",
+    deps = [
+        ":astrowavenet_model",
+        ":configurations",
+        "//astronet/util:configdict",
+    ],
+)

research/astronet/astrowavenet/README.md

+# AstroWaveNet: A generative model for light curves.
+
+Implementation based on "WaveNet: A Generative Model of Raw Audio":
+https://arxiv.org/abs/1609.03499
+
+## Code Authors
+
+Alex Tamkin: [@atamkin](https://github.com/atamkin)
+
+Chris Shallue: [@cshallue](https://github.com/cshallue)
+
+## Pull Requests / Issues
+
+Chris Shallue: [@cshallue](https://github.com/cshallue)
+
+## Additional Dependencies
+
+This package requires TensorFlow 1.12 or greater. As of October 2018, this
+requires the **TensorFlow nightly build**
+([instructions](https://www.tensorflow.org/install/pip)).
+
+In addition to the dependencies listed in the top-level README, this package
+requires:
+
+* **TensorFlow Probability** ([instructions](https://www.tensorflow.org/probability/install))
+* **Six** ([instructions](https://pypi.org/project/six/))
+
+## Basic Usage
+
+To train a model on synthetic transits:
+
+```bash
+bazel build astrowavenet/...
+```
+
+```bash
+bazel-bin/astrowavenet/trainer \
+--dataset=synthetic_transits \
+--model_dir=/tmp/astrowavenet/ \
+--config_overrides='{"hparams": {"batch_size": 16, "num_residual_blocks": 2}}' \
+--schedule=train_and_eval \
+--eval_steps=100 \
+--save_checkpoints_steps=1000
+```

research/astronet/astrowavenet/__init__.py

+# Copyright 2018 The TensorFlow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+

research/astronet/astrowavenet/astrowavenet_model.py

+# Copyright 2018 The TensorFlow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A TensorFlow WaveNet model for generative modeling of light curves.
+
+Implementation based on "WaveNet: A Generative Model of Raw Audio":
+https://arxiv.org/abs/1609.03499
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+
+def _shift_right(x):
+  """Shifts the input Tensor right by one index along the second dimension.
+
+  Pads the front with zeros and discards the last element.
+
+  Args:
+    x: Input three-dimensional tf.Tensor.
+
+  Returns:
+    Padded, shifted tensor of same shape as input.
+  """
+  x_padded = tf.pad(x, [[0, 0], [1, 0], [0, 0]])
+  return x_padded[:, :-1, :]
+
+
+class AstroWaveNet(object):
+  """A TensorFlow model for generative modeling of light curves."""
+
+  def __init__(self, features, hparams, mode):
+    """Basic setup.
+
+    The actual TensorFlow graph is constructed in build().
+
+    Args:
+      features: A dictionary containing "autoregressive_input" and
+        "conditioning_stack", each of which is a named input Tensor. All
+        features have dtype float32 and shape [batch_size, length, dim].
+      hparams: A ConfigDict of hyperparameters for building the model.
+      mode: A tf.estimator.ModeKeys to specify whether the graph should be built
+        for training, evaluation or prediction.
+
+    Raises:
+      ValueError: If mode is invalid.
+    """
+    valid_modes = [
+        tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL,
+        tf.estimator.ModeKeys.PREDICT
+    ]
+    if mode not in valid_modes:
+      raise ValueError("Expected mode in {}. Got: {}".format(valid_modes, mode))
+    self.hparams = hparams
+    self.mode = mode
+
+    self.autoregressive_input = features["autoregressive_input"]
+    self.conditioning_stack = features["conditioning_stack"]
+    self.weights = features.get("weights")
+
+    self.network_output = None  # Sum of skip connections from dilation stack.
+    self.dist_params = None  # Dict of predicted distribution parameters.
+    self.predicted_distributions = None  # Predicted distribution for examples.
+    self.autoregressive_target = None  # Autoregressive target predictions.
+    self.batch_losses = None  # Loss for each predicted distribution in batch.
+    self.per_example_loss = None  # Loss for each example in batch.
+    self.num_nonzero_weight_examples = None  # Number of examples in batch.
+    self.total_loss = None  # Overall loss for the batch.
+    self.global_step = None  # Global step Tensor.
+
+  def causal_conv_layer(self, x, output_size, kernel_width, dilation_rate=1):
+    """Applies a dialated causal convolution to the input.
+
+    Args:
+      x: tf.Tensor; Input tensor.
+      output_size: int; Number of output filters for the convolution.
+      kernel_width: int; Width of the 1D convolution window.
+      dilation_rate: int; Dilation rate of the layer.
+
+    Returns:
+      Resulting tf.Tensor after applying the convolution.
+    """
+    causal_conv_op = tf.keras.layers.Conv1D(
+        output_size,
+        kernel_width,
+        padding="causal",
+        dilation_rate=dilation_rate,
+        name="causal_conv")
+    return causal_conv_op(x)
+
+  def conv_1x1_layer(self, x, output_size, activation=None):
+    """Applies a 1x1 convolution to the input.
+
+    Args:
+      x: tf.Tensor; Input tensor.
+      output_size: int; Number of output filters for the 1x1 convolution.
+      activation: Activation function to apply (e.g. 'relu').
+
+    Returns:
+      Resulting tf.Tensor after applying the 1x1 convolution.
+    """
+    conv_1x1_op = tf.keras.layers.Conv1D(
+        output_size, 1, activation=activation, name="conv1x1")
+    return conv_1x1_op(x)
+
+  def gated_residual_layer(self, x, dilation_rate):
+    """Creates a gated, dilated convolutional layer with a residual connnection.
+
+    Args:
+      x: tf.Tensor; Input tensor
+      dilation_rate: int; Dilation rate of the layer.
+
+    Returns:
+      skip_connection: tf.Tensor; Skip connection to network_output layer.
+      residual_connection: tf.Tensor; Sum of learned residual and input tensor.
+    """
+    with tf.variable_scope("filter"):
+      x_filter_conv = self.causal_conv_layer(x, x.shape[-1].value,
+                                             self.hparams.dilation_kernel_width,
+                                             dilation_rate)
+      cond_filter_conv = self.conv_1x1_layer(self.conditioning_stack,
+                                             x.shape[-1].value)
+    with tf.variable_scope("gate"):
+      x_gate_conv = self.causal_conv_layer(x, x.shape[-1].value,
+                                           self.hparams.dilation_kernel_width,
+                                           dilation_rate)
+      cond_gate_conv = self.conv_1x1_layer(self.conditioning_stack,
+                                           x.shape[-1].value)
+
+    gated_activation = (
+        tf.tanh(x_filter_conv + cond_filter_conv) *
+        tf.sigmoid(x_gate_conv + cond_gate_conv))
+
+    with tf.variable_scope("residual"):
+      residual = self.conv_1x1_layer(gated_activation, x.shape[-1].value)
+    with tf.variable_scope("skip"):
+      skip_connection = self.conv_1x1_layer(gated_activation,
+                                            self.hparams.skip_output_dim)
+
+    return skip_connection, x + residual
+
+  def build_network(self):
+    """Builds WaveNet network.
+
+    This consists of:
+      1) An initial causal convolution,
+      2) The dialation stack, and
+      3) Summing of skip connections
+
+    The network output can then be used to predict various output distributions.
+
+    Inputs:
+      self.autoregressive_input
+      self.conditioning_stack
+
+    Outputs:
+      self.network_output; tf.Tensor
+    """
+    skip_connections = []
+    x = _shift_right(self.autoregressive_input)
+    with tf.variable_scope("preprocess"):
+      x = self.causal_conv_layer(x, self.hparams.preprocess_output_size,
+                                 self.hparams.preprocess_kernel_width)
+    for i in range(self.hparams.num_residual_blocks):
+      with tf.variable_scope("block_{}".format(i)):
+        for dilation_rate in self.hparams.dilation_rates:
+          with tf.variable_scope("dilation_{}".format(dilation_rate)):
+            skip_connection, x = self.gated_residual_layer(x, dilation_rate)
+            skip_connections.append(skip_connection)
+
+    self.network_output = tf.add_n(skip_connections)
+
+  def dist_params_layer(self, x, outputs_size):
+    """Converts x to the correct shape for populating a distribution object.
+
+    Args:
+      x: A Tensor of shape [batch_size, time_series_length, num_features].
+      outputs_size: The number of parameters needed to specify all the
+        distributions in the output. E.g. 5*3=15 to specify 5 distributions with
+        3 parameters each.
+
+    Returns:
+      The parameters of each distribution, a tensor of shape [batch_size,
+        time_series_length, outputs_size].
+    """
+    with tf.variable_scope("dist_params"):
+      conv_outputs = self.conv_1x1_layer(x, outputs_size)
+    return conv_outputs
+
+  def build_predictions(self):
+    """Predicts output distribution from network outputs.
+
+    Runs the model through:
+      1) ReLU
+      2) 1x1 convolution
+      3) ReLU
+      4) 1x1 convolution
+
+    The result of the last convolution is used as the parameters of the
+    specified output distribution (currently either Categorical or Normal).
+
+    Inputs:
+      self.network_outputs
+
+    Outputs:
+      self.dist_params
+      self.predicted_distributions
+
+    Raises:
+      ValueError: If distribution type is neither 'categorical' nor 'normal'.
+    """
+    with tf.variable_scope("postprocess"):
+      network_output = tf.keras.activations.relu(self.network_output)
+      network_output = self.conv_1x1_layer(
+          network_output,
+          output_size=network_output.shape[-1].value,
+          activation="relu")
+    num_dists = self.autoregressive_input.shape[-1].value
+
+    if self.hparams.output_distribution.type == "categorical":
+      num_classes = self.hparams.output_distribution.num_classes
+      logits = self.dist_params_layer(network_output, num_dists * num_classes)
+      logits_shape = tf.concat(
+          [tf.shape(network_output)[:-1], [num_dists, num_classes]], 0)
+      logits = tf.reshape(logits, logits_shape)
+      dist = tfp.distributions.Categorical(logits=logits)
+      dist_params = {"logits": logits}
+    elif self.hparams.output_distribution.type == "normal":
+      loc_scale = self.dist_params_layer(network_output, num_dists * 2)
+      loc, scale = tf.split(loc_scale, 2, axis=-1)
+      # Ensure scale is positive.
+      scale = tf.nn.softplus(scale) + self.hparams.output_distribution.min_scale
+      dist = tfp.distributions.Normal(loc, scale)
+      dist_params = {"loc": loc, "scale": scale}
+    else:
+      raise ValueError("Unsupported distribution type {}".format(
+          self.hparams.output_distribution.type))
+
+    self.dist_params = dist_params
+    self.predicted_distributions = dist
+
+  def build_losses(self):
+    """Builds the training losses.
+
+    Inputs:
+      self.predicted_distributions
+
+    Outputs:
+      self.batch_losses
+      self.total_loss
+    """
+    autoregressive_target = self.autoregressive_input
+
+    # Quantize the target if the output distribution is categorical.
+    if self.hparams.output_distribution.type == "categorical":
+      min_val = self.hparams.output_distribution.min_quantization_value
+      max_val = self.hparams.output_distribution.max_quantization_value
+      num_classes = self.hparams.output_distribution.num_classes
+      clipped_target = tf.keras.backend.clip(autoregressive_target, min_val,
+                                             max_val)
+      quantized_target = tf.floor(
+          (clipped_target - min_val) / (max_val - min_val) * num_classes)
+      # Deal with the corner case where clipped_target equals max_val by mapping
+      # the label num_classes to num_classes - 1. Essentially, this makes the
+      # final quantized bucket a closed interval while all the other quantized
+      # buckets are half-open intervals.
+      quantized_target = tf.where(
+          quantized_target >= num_classes,
+          tf.ones_like(quantized_target) * (num_classes - 1), quantized_target)
+      autoregressive_target = quantized_target
+
+    log_prob = self.predicted_distributions.log_prob(autoregressive_target)
+
+    weights = self.weights
+    if weights is None:
+      weights = tf.ones_like(log_prob)
+    weights_dim = len(weights.shape)
+    per_example_weight = tf.reduce_sum(
+        weights, axis=list(range(1, weights_dim)))
+    per_example_indicator = tf.to_float(tf.greater(per_example_weight, 0))
+    num_examples = tf.reduce_sum(per_example_indicator)
+
+    batch_losses = -log_prob * weights
+    losses_ndims = batch_losses.shape.ndims
+    per_example_loss_sum = tf.reduce_sum(
+        batch_losses, axis=list(range(1, losses_ndims)))
+    per_example_loss = tf.where(per_example_weight > 0,
+                                per_example_loss_sum / per_example_weight,
+                                tf.zeros_like(per_example_weight))
+    total_loss = tf.reduce_sum(per_example_loss) / num_examples
+
+    self.autoregressive_target = autoregressive_target
+    self.batch_losses = batch_losses
+    self.per_example_loss = per_example_loss
+    self.num_nonzero_weight_examples = num_examples
+    self.total_loss = total_loss
+
+  def build(self):
+    """Creates all ops for training, evaluation or inference."""
+    self.global_step = tf.train.get_or_create_global_step()
+
+    self.build_network()
+    self.build_predictions()
+
+    if self.mode in [tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL]:
+      self.build_losses()

research/astronet/astrowavenet/configurations.py

+# Copyright 2018 The TensorFlow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Configurations for model building, training and evaluation."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+def base():
+  """Returns the base config for model building, training and evaluation."""
+  return {
+      # Hyperparameters for building and training the model.
+      "hparams": {
+          "batch_size": 64,
+          "dilation_kernel_width": 2,
+          "skip_output_dim": 10,
+          "preprocess_output_size": 3,
+          "preprocess_kernel_width": 10,
+          "num_residual_blocks": 4,
+          "dilation_rates": [1, 2, 4, 8, 16],
+          "output_distribution": {
+              "type": "normal",
+              "min_scale": 0.001
+          },
+          # Learning rate parameters.
+          "learning_rate": 1e-6,
+          "learning_rate_decay_steps": 0,
+          "learning_rate_decay_factor": 0,
+          "learning_rate_decay_staircase": True,
+
+          # Optimizer for training the model.
+          "optimizer": "adam",
+
+          # If not None, gradient norms will be clipped to this value.
+          "clip_gradient_norm": 1,
+      }
+  }
+
+
+def categorical():
+  """Returns a config for models with a categorical output distribution.
+
+  Input values will be clipped to {min,max}_value_for_quantization, then
+  linearly split into num_classes.
+  """
+  config = base()
+  config["hparams"]["output_distribution"] = {
+      "type": "categorical",
+      "num_classes": 256,
+      "min_quantization_value": -1,
+      "max_quantization_value": 1
+  }
+  return config
+
+
+def get_config(config_name):
+  """Returns config correspnding to provided name."""
+  if config_name in ["base", "normal"]:
+    return base()
+  elif config_name == "categorical":
+    return categorical()
+  else:
+    raise ValueError("Unrecognized config name: {}".format(config_name))

research/astronet/astrowavenet/data/BUILD

+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])  # Apache 2.0
+
+py_library(
+    name = "base",
+    srcs = [
+        "base.py",
+    ],
+    deps = [
+        "//astronet/ops:dataset_ops",
+        "//astronet/util:configdict",
+    ],
+)
+
+py_test(
+    name = "base_test",
+    srcs = ["base_test.py"],
+    data = ["test_data/test-dataset.tfrecord"],
+    srcs_version = "PY2AND3",
+    deps = [":base"],
+)
+
+py_library(
+    name = "kepler_light_curves",
+    srcs = [
+        "kepler_light_curves.py",
+    ],
+    deps = [
+        ":base",
+        "//astronet/util:configdict",
+    ],
+)
+
+py_library(
+    name = "synthetic_transits",
+    srcs = [
+        "synthetic_transits.py",
+    ],
+    deps = [
+        ":base",
+        ":synthetic_transit_maker",
+        "//astronet/util:configdict",
+    ],
+)
+
+py_library(
+    name = "synthetic_transit_maker",
+    srcs = [
+        "synthetic_transit_maker.py",
+    ],
+)
+
+py_test(
+    name = "synthetic_transit_maker_test",
+    srcs = ["synthetic_transit_maker_test.py"],
+    srcs_version = "PY2AND3",
+    deps = [":synthetic_transit_maker"],
+)

research/astronet/astrowavenet/data/__init__.py

+# Copyright 2018 The TensorFlow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+

research/astronet/astrowavenet/data/base.py

+# Copyright 2018 The TensorFlow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Base dataset builder classes for AstroWaveNet input pipelines."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import abc
+import six
+
+import tensorflow as tf
+
+from astronet.util import configdict
+from astronet.ops import dataset_ops
+
+
+@six.add_metaclass(abc.ABCMeta)
+class DatasetBuilder(object):
+  """Base class for building a dataset input pipeline for AstroWaveNet."""
+
+  def __init__(self, config_overrides=None):
+    """Initializes the dataset builder.
+
+    Args:
+      config_overrides: Dict or ConfigDict containing overrides to the default
+        configuration.
+    """
+    self.config = configdict.ConfigDict(self.default_config())
+    if config_overrides is not None:
+      self.config.update(config_overrides)
+
+  @staticmethod
+  def default_config():
+    """Returns the default configuration as a ConfigDict or Python dict."""
+    return {}
+
+  @abc.abstractmethod
+  def build(self, batch_size):
+    """Builds the dataset input pipeline.
+
+    Args:
+      batch_size: The number of input examples in each batch.
+
+    Returns:
+      A tf.data.Dataset object.
+    """
+    raise NotImplementedError
+
+
+@six.add_metaclass(abc.ABCMeta)
+class _ShardedDatasetBuilder(DatasetBuilder):
+  """Abstract base class for a dataset consisting of sharded files."""
+
+  def __init__(self, file_pattern, mode, config_overrides=None, use_tpu=False):
+    """Initializes the dataset builder.
+
+    Args:
+      file_pattern: File pattern matching input file shards, e.g.
+        "/tmp/train-?????-of-00100". May also be a comma-separated list of file
+        patterns.
+      mode: A tf.estimator.ModeKeys.
+      config_overrides: Dict or ConfigDict containing overrides to the default
+        configuration.
+      use_tpu: Whether to build the dataset for TPU.
+    """
+    super(_ShardedDatasetBuilder, self).__init__(config_overrides)
+    self.file_pattern = file_pattern
+    self.mode = mode
+    self.use_tpu = use_tpu
+
+  @staticmethod
+  def default_config():
+    config = super(_ShardedDatasetBuilder,
+                   _ShardedDatasetBuilder).default_config()
+    config.update({
+        "max_length": 1024,
+        "shuffle_values_buffer": 1000,
+        "num_parallel_parser_calls": 4,
+        "batches_buffer_size": None,  # Defaults to max(1, 256 / batch_size).
+    })
+    return config
+
+  @abc.abstractmethod
+  def file_reader(self):
+    """Returns a function that reads a single sharded file."""
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def create_example_parser(self):
+    """Returns a function that parses a single tf.Example proto."""
+    raise NotImplementedError
+
+  def _batch_and_pad(self, dataset, batch_size):
+    """Combines elements into batches of the same length, padding if needed."""
+    if self.use_tpu:
+      padded_length = self.config.max_length
+      if not padded_length:
+        raise ValueError("config.max_length is required when using TPU")
+      # Pad with zeros up to padded_length. Note that this will pad the
+      # "weights" Tensor with zeros as well, which ensures that padded elements
+      # do not contribute to the loss.
+      padded_shapes = {}
+      for name, shape in dataset.output_shapes.iteritems():
+        shape.assert_is_compatible_with([None, None])  # Expect a 2D sequence.
+        dims = shape.as_list()
+        dims[0] = padded_length
+        shape = tf.TensorShape(dims)
+        shape.assert_is_fully_defined()
+        padded_shapes[name] = shape
+    else:
+      # Pad each batch up to the maximum size of each dimension in the batch.
+      padded_shapes = dataset.output_shapes
+
+    return dataset.padded_batch(batch_size, padded_shapes)
+
+  def build(self, batch_size):
+    """Builds the dataset input pipeline.
+
+    Args:
+      batch_size:
+
+    Returns:
+      A tf.data.Dataset.
+
+    Raises:
+      ValueError: If no files match self.file_pattern.
+    """
+    file_patterns = self.file_pattern.split(",")
+    filenames = []
+    for p in file_patterns:
+      matches = tf.gfile.Glob(p)
+      if not matches:
+        raise ValueError("Found no input files matching {}".format(p))
+      filenames.extend(matches)
+    tf.logging.info(
+        "Building input pipeline from %d files matching patterns: %s",
+        len(filenames), file_patterns)
+
+    is_training = self.mode == tf.estimator.ModeKeys.TRAIN
+
+    # Create a string dataset of filenames, and possibly shuffle.
+    filename_dataset = tf.data.Dataset.from_tensor_slices(filenames)
+    if is_training and len(filenames) > 1:
+      filename_dataset = filename_dataset.shuffle(len(filenames))
+
+    # Read serialized Example protos.
+    dataset = filename_dataset.apply(
+        tf.contrib.data.parallel_interleave(
+            self.file_reader(), cycle_length=8, block_length=8, sloppy=True))
+
+    if is_training:
+      # Shuffle and repeat. Note that shuffle() is before repeat(), so elements
+      # are shuffled among each epoch of data, and not between epochs of data.
+      if self.config.shuffle_values_buffer > 0:
+        dataset = dataset.shuffle(self.config.shuffle_values_buffer)
+      dataset = dataset.repeat()
+
+    # Map the parser over the dataset.
+    dataset = dataset.map(
+        self.create_example_parser(),
+        num_parallel_calls=self.config.num_parallel_parser_calls)
+
+    def _prepare_wavenet_inputs(features):
+      """Validates features, and clips lengths and adds weights if needed."""
+      # Validate feature names.
+      required_features = {"autoregressive_input", "conditioning_stack"}
+      allowed_features = required_features | {"weights"}
+      feature_names = features.keys()
+      if not required_features.issubset(feature_names):
+        raise ValueError("Features must contain all of: {}. Got: {}".format(
+            required_features, feature_names))
+      if not allowed_features.issuperset(feature_names):
+        raise ValueError("Features can only contain: {}. Got: {}".format(
+            allowed_features, feature_names))
+
+      output = {}
+      for name, value in features.items():
+        # Validate shapes. The output dimension is [num_samples, dim].
+        ndims = len(value.shape)
+        if ndims == 1:
+          # Add an extra dimension: [num_samples] -> [num_samples, 1].
+          value = tf.expand_dims(value, -1)
+        elif ndims != 2:
+          raise ValueError(
+              "Features should be 1D or 2D sequences. Got '{}' = {}".format(
+                  name, value))
+        if self.config.max_length:
+          value = value[:self.config.max_length]
+        output[name] = value
+
+      if "weights" not in output:
+        output["weights"] = tf.ones_like(output["autoregressive_input"])
+
+      return output
+
+    dataset = dataset.map(_prepare_wavenet_inputs)
+
+    # Batch results by up to batch_size.
+    dataset = self._batch_and_pad(dataset, batch_size)
+
+    if is_training:
+      # The dataset repeats infinitely before batching, so each batch has the
+      # maximum number of elements.
+      dataset = dataset_ops.set_batch_size(dataset, batch_size)
+    elif self.use_tpu and self.mode == tf.estimator.ModeKeys.EVAL:
+      # Pad to ensure that each batch has the same number of elements.
+      dataset = dataset_ops.pad_dataset_to_batch_size(dataset, batch_size)
+
+    # Prefetch batches.
+    buffer_size = (
+        self.config.batches_buffer_size or max(1, int(256 / batch_size)))
+    dataset = dataset.prefetch(buffer_size)
+
+    return dataset
+
+
+def tfrecord_reader(filename):
+  """Returns a tf.data.Dataset that reads a single TFRecord file shard."""
+  return tf.data.TFRecordDataset(filename, buffer_size=16 * 1000 * 1000)
+
+
+class TFRecordDataset(_ShardedDatasetBuilder):
+  """Builder for a dataset consisting of TFRecord files."""
+
+  def file_reader(self):
+    """Returns a function that reads a single file shard."""
+    return tfrecord_reader

research/astronet/astrowavenet/data/kepler_light_curves.py

+# Copyright 2018 The TensorFlow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Kepler light curve inputs to the AstroWaveNet model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+from astrowavenet.data import base
+
+COND_INPUT_KEY = "mask"
+AR_INPUT_KEY = "flux"
+
+
+class KeplerLightCurves(base.TFRecordDataset):
+  """Kepler light curve inputs to the AstroWaveNet model."""
+
+  def create_example_parser(self):
+
+    def _example_parser(serialized):
+      """Parses a single tf.Example proto."""
+      features = tf.parse_single_example(
+          serialized,
+          features={
+              AR_INPUT_KEY: tf.VarLenFeature(tf.float32),
+              COND_INPUT_KEY: tf.VarLenFeature(tf.int64),
+          })
+      # Extract values from SparseTensor objects.
+      autoregressive_input = features[AR_INPUT_KEY].values
+      conditioning_stack = tf.to_float(features[COND_INPUT_KEY].values)
+      return {
+          "autoregressive_input": autoregressive_input,
+          "conditioning_stack": conditioning_stack,
+      }
+
+    return _example_parser

research/astronet/astrowavenet/data/synthetic_transit_maker.py

+# Copyright 2018 The TensorFlow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Generates synthetic light curves with periodic transit-like dips.
+
+See class docstring below for more information.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+
+class SyntheticTransitMaker(object):
+  """Generates synthetic light curves with periodic transit-like dips.
+
+  These light curves are generated by thresholding noisy sine waves. Each time
+  random_light_curve is called, a thresholded sine wave is generated by sampling
+  parameters uniformly from the ranges specified below.
+
+  Attributes:
+    period_range: A tuple of positive values specifying the range of periods the
+      sine waves may take.
+    amplitude_range: A tuple of positive values specifying the range of
+      amplitudes the sine waves may take.
+    threshold_ratio_range: A tuple of values in [0, 1) specifying the range of
+      thresholds as a ratio of the sine wave amplitude.
+    phase_range: Tuple of values specifying the range of phases the sine wave
+      may take as a ratio of the sampled period. E.g. a sampled phase of 0.5
+      would translate the sine wave by half of the period. The most common
+      reason to override this would be to generate light curves
+      deterministically (with e.g. (0,0)).
+    noise_sd_range: A tuple of values in [0, 1) specifying the range of standard
+      deviations for the Gaussian noise applied to the sine wave.
+  """
+
+  def __init__(self,
+               period_range=(0.5, 4),
+               amplitude_range=(1, 1),
+               threshold_ratio_range=(0, 0.99),
+               phase_range=(0, 1),
+               noise_sd_range=(0.1, 0.1)):
+
+    if threshold_ratio_range[0] < 0 or threshold_ratio_range[1] >= 1:
+      raise ValueError("Threshold ratio range must be in [0, 1). Got: {}."
+                       .format(threshold_ratio_range))
+    if amplitude_range[0] <= 0:
+      raise ValueError(
+          "Amplitude range must only contain positive numbers. Got: {}.".format(
+              amplitude_range))
+    if period_range[0] <= 0:
+      raise ValueError(
+          "Period range must only contain positive numbers. Got: {}.".format(
+              period_range))
+    if noise_sd_range[0] < 0:
+      raise ValueError(
+          "Noise standard deviation range must be nonnegative. Got: {}.".format(
+              noise_sd_range))
+
+    for (start, end), name in [(period_range, "period"),
+                               (amplitude_range, "amplitude"),
+                               (threshold_ratio_range, "threshold ratio"),
+                               (phase_range, "phase range"),
+                               (noise_sd_range, "noise standard deviation")]:
+      if end < start:
+        raise ValueError(
+            "End of {} range may not be less than start. Got: ({}, {})".format(
+                name, start, end))
+
+    self.period_range = period_range
+    self.amplitude_range = amplitude_range
+    self.threshold_ratio_range = threshold_ratio_range
+    self.phase_range = phase_range
+    self.noise_sd_range = noise_sd_range
+
+  def random_light_curve(self, time, mask_prob=0):
+    """Samples parameters and generates a light curve.
+
+    Args:
+      time: np.array, x-values to sample from the thresholded sine wave.
+      mask_prob: value in [0,1], probability an individual datapoint is set to
+        zero
+
+    Returns:
+      flux: np.array, values of the masked sampled light curve corresponding to
+        the provided time array.
+      mask: np.array of ones and zeros, with zeros indicating masking at the
+        respective position on the flux array.
+    """
+
+    period = np.random.uniform(*self.period_range)
+    phase = np.random.uniform(*self.phase_range) * period
+    amplitude = np.random.uniform(*self.amplitude_range)
+    threshold = np.random.uniform(*self.threshold_ratio_range) * amplitude
+
+    sin_wave = np.sin(time / period - phase) * amplitude
+    flux = np.minimum(sin_wave, -threshold) + threshold
+
+    noise_sd = np.random.uniform(*self.noise_sd_range)
+    noise = np.random.normal(scale=noise_sd, size=(len(time),))
+    flux += noise
+
+    # Array of ones and zeros, where zeros indicate masking.
+    mask = np.random.random(len(time)) > mask_prob
+    mask = mask.astype(np.float)
+
+    return flux * mask, mask
+
+  def random_light_curve_generator(self, time, mask_prob=0):
+    """Returns a generator function yielding random light curves.
+
+    Args:
+       time: An np.array of x-values to sample from the thresholded sine wave.
+       mask_prob: Value in [0,1], probability an individual datapoint is set to
+         zero.
+
+    Returns:
+      A generator yielding random light curves.
+    """
+
+    def generator_fn():
+      while True:
+        yield self.random_light_curve(time, mask_prob)
+
+    return generator_fn

research/astronet/astrowavenet/data/synthetic_transit_maker_test.py

+# Copyright 2018 The TensorFlow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for synthetic_transit_maker."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+import numpy as np
+
+from astrowavenet.data import synthetic_transit_maker
+
+
+class SyntheticTransitMakerTest(absltest.TestCase):
+
+  def testBadRangesRaiseExceptions(self):
+
+    # Period range cannot contain negative values.
+    with self.assertRaisesRegexp(ValueError, "Period"):
+      synthetic_transit_maker.SyntheticTransitMaker(period_range=(-1, 10))
+
+    # Amplitude range cannot contain negative values.
+    with self.assertRaisesRegexp(ValueError, "Amplitude"):
+      synthetic_transit_maker.SyntheticTransitMaker(amplitude_range=(-10, -1))
+
+    # Threshold ratio range must be contained in the half-open interval [0, 1).
+    with self.assertRaisesRegexp(ValueError, "Threshold ratio"):
+      synthetic_transit_maker.SyntheticTransitMaker(
+          threshold_ratio_range=(0, 1))
+
+    # Noise standard deviation range must only contain nonnegative values.
+    with self.assertRaisesRegexp(ValueError, "Noise standard deviation"):
+      synthetic_transit_maker.SyntheticTransitMaker(noise_sd_range=(-1, 1))
+
+    # End of range may not be less than start.
+    invalid_range = (0.2, 0.1)
+    range_args = [
+        "period_range", "threshold_ratio_range", "amplitude_range",
+        "noise_sd_range", "phase_range"
+    ]
+    for range_arg in range_args:
+      with self.assertRaisesRegexp(ValueError, "may not be less"):
+        synthetic_transit_maker.SyntheticTransitMaker(
+            **{range_arg: invalid_range})
+
+  def testStochasticLightCurveGeneration(self):
+    transit_maker = synthetic_transit_maker.SyntheticTransitMaker()
+
+    time = np.arange(100)
+    flux, mask = transit_maker.random_light_curve(time, mask_prob=0.4)
+    self.assertEqual(len(flux), 100)
+    self.assertEqual(len(mask), 100)
+
+  def testDeterministicLightCurveGeneration(self):
+    gold_flux = np.array([
+        0., 0., 0., 0., 0., 0., 0., -0.85099258, -2.04776251, -2.65829632,
+        -2.53014378, -1.69530454, -0.36223792, 0., 0., 0., 0., 0., 0.,
+        -0.2110405, -1.57757635, -2.47528153, -2.67999913, -2.14061117,
+        -0.9918028, 0., 0., 0., 0., 0., 0., 0., -1.01475559, -2.15534176,
+        -2.68282928, -2.46550457, -1.55763357, -0.18591162, 0., 0., 0., 0., 0.,
+        0., -0.3870683, -1.71426199, -2.53849461, -2.65395535, -2.03181367,
+        -0.82741829, 0., 0., 0., 0., 0., 0., 0., -1.17380391, -2.2541162,
+        -2.69666588, -2.39094831, -1.41330116, -0.00784284, 0., 0., 0., 0., 0.,
+        0., -0.56063229, -1.84372452, -2.59152891, -2.61731875, -1.91465433,
+        -0.65899089, 0., 0., 0., 0., 0., 0., 0., -1.3275672, -2.34373163,
+        -2.69975648, -2.30674237, -1.26282489, 0., 0., 0., 0., 0., 0., 0.,
+        -0.73111006, -1.9654997, -2.63419424, -2.5702207, -1.78955328,
+        -0.48712456
+    ])
+
+    # Use ranges containing one value for determinism.
+    transit_maker = synthetic_transit_maker.SyntheticTransitMaker(
+        period_range=(2, 2),
+        amplitude_range=(3, 3),
+        threshold_ratio_range=(.1, .1),
+        phase_range=(0, 0),
+        noise_sd_range=(0, 0))
+
+    time = np.linspace(0, 100, 100)
+
+    flux, mask = transit_maker.random_light_curve(time)
+    np.testing.assert_array_almost_equal(flux, gold_flux)
+    np.testing.assert_array_almost_equal(mask, np.ones(100))
+
+  def testRandomLightCurveGenerator(self):
+    transit_maker = synthetic_transit_maker.SyntheticTransitMaker()
+    time = np.linspace(0, 100, 100)
+    generator = transit_maker.random_light_curve_generator(
+        time, mask_prob=0.3)()
+    for _ in range(5):
+      flux, mask = next(generator)
+      self.assertEqual(len(flux), 100)
+      self.assertEqual(len(mask), 100)
+
+
+if __name__ == "__main__":
+  absltest.main()