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()