modeling.py

# Copyright 2017 Google LLC.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from this
#    software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
"""Provides an abstraction around deep learning models in DeepVariant.

This class allows us to encapsulate all of the model management, loading,
saving, and data processing in a single place so those details don't spill over
into the more general deepvariant codebase. The key thing we are aiming for here
is to make sure we can easily play with other model architectures without
modifying the surrounding training and evaluation code.
"""

import enum
import itertools
import math



from absl import flags
from absl import logging
from tensorflow import estimator as tf_estimator
from tensorflow.compat.v1 import estimator as tf_compat_v1_estimator
import tensorflow as tf
import tf_slim
from tf_slim.nets import inception_v3

from deepvariant import dv_constants
from deepvariant import tf_utils
# pylint: disable=g-direct-tensorflow-import
from tensorflow.python.framework import ops
from tensorflow.python.tpu import tpu_config
from tensorflow.python.tpu import tpu_estimator
from tensorflow.python.tpu import tpu_optimizer
# pylint: enable=g-direct-tensorflow-import
from deepvariant import attention_inception_v3

tf.compat.v1.disable_eager_execution()

flags.DEFINE_float(
    'label_smoothing', 1e-6,
    'Amount of label smoothing to use. By default this is 0.0001% '
    'meaning that we expect a label error at a rate of 1 / 1,000,000')

# Training parameters.
flags.DEFINE_float('learning_rate', 0.064, 'Initial learning rate.')

flags.DEFINE_float('rmsprop_momentum', 0.9, 'Momentum.')

flags.DEFINE_float('rmsprop_decay', 0.9, 'Decay term for RMSProp.')

flags.DEFINE_float('rmsprop_epsilon', 1.0, 'Epsilon term for RMSProp.')

flags.DEFINE_float('learning_rate_decay_factor', 0.94,
                   'Learning rate decay factor.')

flags.DEFINE_float('num_epochs_per_decay', 2.0,
                   'Number of epochs after which learning rate decays.')

flags.DEFINE_float('moving_average_decay', 0.9999,
                   'The decay to use for the moving average.')

flags.DEFINE_integer(
    'save_summary_steps', 100, 'Number of steps which must have run before '
    'showing summaries.')

flags.DEFINE_integer(
    'save_interval_secs', 60 * 10,
    'Interval (in seconds) at which the model data '
    'should be checkpointed. Set to 0 to disable, -1 to ignore. '
    'Exclusive with save_interval_steps.')

flags.DEFINE_integer(
    'save_interval_steps', -1, 'Interval (in steps) at which the model data '
    'should be checkpointed. Set to 0 to disable, -1 to ignore. '
    'Exclusive with save_interval_secs.')

flags.DEFINE_integer(
    'seq_type_embedding_size', 200,
    'Set the embedding size for the sequencing type embeddings. Default is 200. '
    'This flag is only useful when model_name is `inception_v3_embedding`.')

flags.DEFINE_bool(
    'allow_warmstart_from_different_num_channels', False,
    'If True, always allow warmstarting from model checkpoints '
    'that has different number of channels.')

FLAGS = flags.FLAGS

slim = tf_slim


class UnsupportedImageDimensionsError(Exception):
  """Exception indicating the image dimensions aren't supported by our model."""


def binarize(labels, target_class):
  """Binarize labels and predictions.

  The labels that are equal to target_class parameter are set to 0, else
  set to 1.

  Args:
    labels: the ground-truth labels for the examples.
    target_class: index of the class that is left as zero.

  Returns:
    Tensor of the same shape as labels.
  """
  labels_binary = tf.compat.v1.where(
      tf.equal(labels, tf.constant(target_class, dtype=tf.int64)),
      tf.zeros_like(labels), tf.ones_like(labels))
  return labels_binary


def get_class_recall(labels, predicted_class, target_class):
  """Compute recall from labels and predicted_class for target_class.

  Examples with label target_class are positives. Other classes are negatives.

  Args:
    labels: the ground-truth labels for the examples.
    predicted_class: the predicted labels for the examples.
    target_class: index of the class that is left as non-zero.

  Returns:
    Tensor containing the recall value.
  """
  labels_binary = binarize(labels, target_class)
  predicted_class_binary = binarize(predicted_class, target_class)
  return tf.compat.v1.metrics.recall(labels_binary, predicted_class_binary)


def get_class_precision(labels, predicted_class, target_class):
  """Compute precision from labels and predicted_class for target_class.

  Examples with label target_class are positives. Other classes are negatives.

  Args:
    labels: the ground-truth labels for the examples.
    predicted_class: the predicted labels for the examples.
    target_class: index of the class that is left as non-zero.

  Returns:
    Tensor containing the precision value.
  """
  labels_binary = binarize(labels, target_class)
  predicted_class_binary = binarize(predicted_class, target_class)
  return tf.compat.v1.metrics.precision(labels_binary, predicted_class_binary)


# TODO: Verify this F1 score is correct.
def get_f1_score(labels, predictions, target_class=None):
  """Compute F1 score of predictions with respect to the labels.

  Args:
    labels: tensor whose dimensions must match predictions. The ground-truth
      labels for the examples.
    predictions: tensor of arbitrary dimension. The predicted labels for the
      examples.
    target_class: int. Index of the class that is left as non-zero.

  Returns:
    f1_score: scalar float tensor whose dimensions match predictions. The
    calculated f1 score.
    update_op: operation that updates the f1 score streaming metric.
  """
  if target_class:
    labels = binarize(labels, target_class)
    predictions = binarize(predictions, target_class)

  precision, precision_op = tf.compat.v1.metrics.precision(labels, predictions)
  recall, recall_op = tf.compat.v1.metrics.recall(labels, predictions)

  def compute_f1_score(name):
    pr_product = tf.multiply(precision, recall)
    return tf.math.divide(
        tf.multiply(2.0, pr_product), tf.add(tf.add(precision, recall), 1e-12),
        name)

  f1_score = compute_f1_score('value')
  with ops.control_dependencies([precision_op, recall_op]):
    update_op = compute_f1_score('update_op')

  return f1_score, update_op


def is_encoded_variant_type(variant_types_tensor, type_to_select):
  """Returns a bool tensor indicating which variant_types match type_to_select.

  Args:
    variant_types_tensor: Tensor of shape (batch_size, 1) containing
      EncodedVariantType.value int64 values. Each element of this tensor should
      be a EncodedVariantType.value int64 value indicating the type of the
      variant.
    type_to_select: EncodedVariantType. The type of variant we want to select.

  Returns:
    Tensor of shape (batch_size, 1) of type tf.bool. A True value indicates that
    the variant_type at that position matched type_to_select. Has a False
    otherwise.
  """
  return tf.equal(variant_types_tensor,
                  tf.constant(type_to_select.value, dtype=tf.int64))


# This dictionary contains a mapping from the human readable name of a metric
# function (e.g., Accuracy) and its associated TensorFlow metric function. All
# of the entries here will be stratified by variant_type in eval_metric_fn.
_METRICS_FUNCS_BY_VARIANT_TYPE = {
    'Accuracy': tf.compat.v1.metrics.accuracy,
    'Precision': tf.compat.v1.metrics.precision,
    'Recall': tf.compat.v1.metrics.recall,
    'FPs': tf.compat.v1.metrics.false_positives,
    'FNs': tf.compat.v1.metrics.false_negatives,
    'TPs': tf.compat.v1.metrics.true_positives,
    'TNs': tf.compat.v1.metrics.true_negatives,
}

# A set containing the names of the variant types we split our metrics by type
# by. This data structure isn't a dictionary like it's neighbors because
# eval_metric_fn requires special logic to compute the values here associated
# with each of these names.
_METRICS_BY_VARIANT_TYPE = {'All', 'SNPs', 'Indels'}

# This dictionary contains a mapping from the human readable name of a genotype
# class (e.g., Het) and its associated class label (e.g., 1). All of the entries
# here will be stratified by genotype_class in eval_metric_fn.
_METRICS_GENOTYPE_CLASSES = {
    'HomRef': 0,
    'Het': 1,
    'HomVar': 2,
}

# This dictionary contains a mapping from the human readable name of a metric
# function (e.g., Accuracy) and its associated metric function. All
# of the entries here will be stratified by genotype class (e.g., Het) in
# eval_metric_fn.
_METRICS_FUNCS_BY_GENOTYPE_CLASS = {
    'Precision': get_class_precision,
    'Recall': get_class_recall,
    'F1': get_f1_score,
}


def _eval_name(metric_name, stratification_name):
  return metric_name + '/' + stratification_name


class EvalMetricOrdering(enum.Enum):
  """Enum capturing whether a better metric should be larger or smaller."""
  BIGGER_IS_BETTER = 1
  SMALLER_IS_BETTER = 2


def eval_function_metrics(has_variant_types=True):
  """Gets the set of eval_metrics names and their directionality.

  Args:
    has_variant_types: bool. Will we be providing variant_type information
      during eval so that we'll have metrics stratified by variant_type?

  Returns:
    dict mapping from a metric name string (e.g., "F1/All") and a
    EvalMetricOrdering enum indicating whether larger metric values are better
    or worse.
  """
  names = {_eval_name('F1', 'All'): EvalMetricOrdering.BIGGER_IS_BETTER}

  if has_variant_types:
    variant_type_names = _METRICS_BY_VARIANT_TYPE
  else:
    variant_type_names = {'All'}
  for m, s in itertools.product(_METRICS_FUNCS_BY_VARIANT_TYPE,
                                variant_type_names):
    names[_eval_name(m, s)] = EvalMetricOrdering.BIGGER_IS_BETTER

  for m, s in itertools.product(_METRICS_FUNCS_BY_GENOTYPE_CLASS,
                                _METRICS_GENOTYPE_CLASSES):
    names[_eval_name(m, s)] = EvalMetricOrdering.BIGGER_IS_BETTER

  return names


# NB. This includes only a subset of our usual metrics.
# We'll add the rest back in a subsequent change.
def eval_metric_fn(labels, predictions, variant_types):
  """Calculate eval metrics from Tensors, on CPU host.

  Args:
    labels: the ground-truth labels for the examples.
    predictions: the predicted labels for the examples.
    variant_types: variant types (int64 of EncodedVariantType.value) as a tensor
      of (batch_size,) or None. The types of these variants. If None, no type
      specific evals will be performed.

  Returns:
    A dictionary of string name to metric.
  """
  predicted_classes = tf.argmax(input=predictions, axis=1)

  metrics = {}

  # Add the metrics stratified by variant_type
  weights_by_type = {'All': None}
  if variant_types is not None:
    weights_by_type['SNPs'] = is_encoded_variant_type(
        variant_types, tf_utils.EncodedVariantType.SNP)
    weights_by_type['Indels'] = is_encoded_variant_type(
        variant_types, tf_utils.EncodedVariantType.INDEL)

  for metric_name, metric_func in _METRICS_FUNCS_BY_VARIANT_TYPE.items():
    for weight_name, weights in weights_by_type.items():
      metrics[_eval_name(metric_name, weight_name)] = metric_func(
          labels, predicted_classes, weights=weights)

  # Add the metrics stratified by predicted class.
  for metric_name, metric_func in _METRICS_FUNCS_BY_GENOTYPE_CLASS.items():
    for class_name, class_value in _METRICS_GENOTYPE_CLASSES.items():
      metrics[_eval_name(metric_name,
                         class_name)] = metric_func(labels, predicted_classes,
                                                    class_value)

  # Special case F1/All to avoid a clash between the two different ways that we
  # can compute Precision and Recall (e.g., get_class_precision vs.
  # tf.compat.v1.metrics.precision.
  metrics[_eval_name('F1', 'All')] = get_f1_score(labels, predicted_classes)

  logging.info('Metrics are %s', metrics.keys())

  # Make sure our metrics are consistent with the expected names from
  # eval_function_metrics.
  expected_metrics = eval_function_metrics(
      has_variant_types=variant_types is not None)
  if set(expected_metrics) != set(metrics):
    raise AssertionError(
        'Bug: actual metrics={} not equal to expected={}'.format(
            ','.join(metrics), ','.join(expected_metrics)))

  return metrics


# The following two classes support loading exponential moving averages into
# their corresponding variables when a checkpoint is loaded. They're called
# as hooks by the Estimators. Note for future work: this is the documented
# way, but someone on the mailing list suggested that using the scaffold_fn
# mechanism might be better.


class LoadEMAHook(tf_estimator.SessionRunHook):
  """Hook to load EMA into their corresponding variables.

  This looks for the latest checkpoint in the model dir.
  """

  def __init__(self, model_dir, ignore_missing_vars=False):
    super(LoadEMAHook, self).__init__()
    self._model_dir = model_dir
    self._ignore_missing_vars = ignore_missing_vars

  def begin(self):
    ema = tf.train.ExponentialMovingAverage(FLAGS.moving_average_decay)
    variables_to_restore = ema.variables_to_restore()
    self._load_ema = slim.assign_from_checkpoint_fn(
        tf.train.latest_checkpoint(self._model_dir),
        variables_to_restore,
        ignore_missing_vars=self._ignore_missing_vars)

  def after_create_session(self, sess, coord):
    tf.compat.v1.logging.info('Reloading EMA...')
    self._load_ema(sess)


class PredictEMAHook(tf_estimator.SessionRunHook):
  """Hook to load EMA into their corresponding variables.

  This reads the specified checkpoint.
  """

  def __init__(self, checkpoint_path, ignore_missing_vars=False):
    super(PredictEMAHook, self).__init__()
    self._checkpoint_path = checkpoint_path
    self._ignore_missing_vars = ignore_missing_vars

  def begin(self):
    ema = tf.train.ExponentialMovingAverage(FLAGS.moving_average_decay)
    variables_to_restore = ema.variables_to_restore()
    self._load_ema = slim.assign_from_checkpoint_fn(
        self._checkpoint_path,
        variables_to_restore,
        ignore_missing_vars=self._ignore_missing_vars)

  def after_create_session(self, sess, coord):
    tf.compat.v1.logging.info('Reloading EMA...')
    self._load_ema(sess)


class DeepVariantModel(object):
  """Base class for models that compute genotype likelihoods from an image.

  This class is intended for use anywhere in DeepVariant where we want to train
  or evaluate a model that computes genotype likelihoods from a pileup image. A
  bit of encapsulation helps us to try new models (beyond inception_v3) and unit
  test our code.

  The base class cannot be used directly; concrete subclasses actually implement
  specific models and all of the associated machinery to create/load/save
  models.

  Attributes:
    name: str. The name of this model, such as `inception_v3`.
    pretrained_model_path: str. Path to a root checkpoint where we can start
      training the model, if we are not starting from scratch.
    supported_dimensions_message: str. A human-readable string containing info
      about what image dimensions are supported by this model. E.g., "only
      widths between 42 and 189".
    use_tpu: bool or None. If True, we are executing the model on a TPU, False
      if we are using some other hardware. If None, the execution hardware is
      not yet known.
    model_dir: str or None. The path to the location where model checkpoint are
      being stored. If None, the path hasn't been set yet or is unknown.
  """

  def __init__(self, name, pretrained_model_path):
    """Creates a new DeepVariantModel with name and pretrained_model_path.

    Args:
      name: str. The name of the model. Passed to DeepVariantModel name.
      pretrained_model_path: str. A path to a pretrained model to initialize our
        network from when starting from the 'model_default'.  If None, training
        will start from randomly-initialized parameters.

    Raises:
      ValueError: if any of the arguments is invalid.
    """
    if not name:
      raise ValueError('Got an empty value for name', name)
    self.name = name
    self.pretrained_model_path = pretrained_model_path
    self.supported_dimensions_message = 'unknown'
    self.use_tpu = None

    # Set the model_dir to None by default. We capture its actual value during
    # a call to make_estimator below.
    self.model_dir = None

  def construct_scalar_host_call(self,
                                 metric_dict,
                                 model_dir,
                                 prefix='',
                                 record_frequency_in_steps=100):
    """Construct a host call to log scalars when training on TPU.

    Args:
      metric_dict: A dict of the tensors to be logged.
      model_dir: The location to write the summary.
      prefix: The prefix (if any) to prepend to the metric names.
      record_frequency_in_steps: int; How often should we log our metrics in
        step units.

    Returns:
      A tuple of (function, args_to_be_passed_to_said_function)
    """
    # type: (dict, str) -> (function, list)
    metric_names = list(metric_dict.keys())

    def host_call_fn(global_step, *args):
      """Training host call.

      Creates scalar summaries for training metrics.

      This function is executed on the CPU and should not directly reference
      any Tensors in the rest of the `model_fn`. To pass Tensors from the
      model to the `metric_fn`, provide as part of the `host_call`. See
      https://www.tensorflow.org/api_docs/python/tf/compat/v1/estimator/tpu/TPUEstimator
      for more information.
      Arguments should match the list of `Tensor` objects passed as the second
      element in the tuple passed to `host_call`.
      Args:
        global_step: Tensor with shape `[batch]` for the global_step
        *args: Remaining tensors to log.

      Returns:
        List of summary ops to run on the CPU host.
      """
      step = global_step[0]
      with tf.compat.v2.summary.create_file_writer(
          logdir=model_dir, filename_suffix='.host_call').as_default():
        with tf.compat.v2.summary.record_if(
            lambda: tf.math.equal(step % record_frequency_in_steps, 0)):
          for i, name in enumerate(metric_names):
            tf.compat.v2.summary.scalar(
                name=prefix + name, data=args[i][0], step=step)
          return tf.compat.v1.summary.all_v2_summary_ops()

    # To log the current learning rate, and gradient norm for Tensorboard, the
    # summary op needs to be run on the host CPU via host_call. host_call
    # expects [batch_size, ...] Tensors, thus reshape to introduce a batch
    # dimension. These Tensors are implicitly concatenated to
    # [params['batch_size']].
    global_step_tensor = tf.reshape(
        tf.compat.v1.train.get_or_create_global_step(), [1])
    other_tensors = [tf.reshape(metric_dict[key], [1]) for key in metric_names]

    return host_call_fn, [global_step_tensor] + other_tensors

  def _create_warm_start_settings(self, start_from_checkpoint):
    """Create a proper WarmStartSettings based on start_from_checkpoint."""
    # If the special value "model_default" was passed, ask the model for
    # its default.
    if start_from_checkpoint == 'model_default':
      start_from_checkpoint = self.pretrained_model_path

    # If the path is non-False, use it.
    if start_from_checkpoint:
      logging.info('Initializing model from checkpoint at %s',
                   start_from_checkpoint)
      excluded_scopes = set()
      reader = tf.compat.v1.train.NewCheckpointReader(start_from_checkpoint)
      var_to_shape_map = reader.get_variable_to_shape_map()
      if tf_utils.model_num_classes(
          start_from_checkpoint,
          self.n_classes_model_variable) != dv_constants.NUM_CLASSES:
        excluded_scopes.update(self.excluded_scopes_for_incompatible_classes)
      if FLAGS.allow_warmstart_from_different_num_channels:
        excluded_scopes.update(self.excluded_scopes_for_incompatible_channels)
      if excluded_scopes:
        logging.info(
            'The model checkpoint to warm start from has different '
            'shapes. If this is in training, we will '
            'exclude: %s', excluded_scopes)
        vars_to_include = [
            v for v in var_to_shape_map.keys()
            if not v.startswith(tuple(excluded_scopes))
        ]
      else:
        logging.info('The model checkpoint to warm start from should have the '
                     'same number of classes and same numbers of channels.'
                     'If this is in training, we will include everything for '
                     'warm starting....')
        vars_to_include = var_to_shape_map.keys()
      return tf_estimator.WarmStartSettings(
          ckpt_to_initialize_from=start_from_checkpoint,
          vars_to_warm_start='|'.join(vars_to_include))
    else:
      # If warm_start_from is an empty string, specifically set it to None.
      logging.vlog(3, 'Initializing model with random parameters')
      return None

  def make_estimator(self,
                     batch_size,
                     model_dir=None,
                     max_checkpoints_to_keep=100000,
                     iterations_per_loop=100,
                     params=None,
                     unused_device_fn=None,
                     master='',
                     use_tpu=False,
                     start_from_checkpoint=None,
                     session_config=None,
                     include_debug_info=False):
    """Returns a new tf.estimator.Estimator object for training or prediction.

    The estimator needs to know batch_size. We use the same value for all
    of eval, train, and predict. The estimator will automatically save
    checkpoints to model_dir and keep the specified number of them. The value
    of iterations_per_loop is not critical, and we default to the recommended
    value. Some optional arguments are only required for use with TPU.

    This function will use self.model_fn and self.use_tpu when constructing the
    model specific Estimator object.

    Estimators are also sometimes called classifiers.

    Args:
      batch_size: the batch size to use (for TRAIN, EVAL, and PREDICT modes).
      model_dir: an (optional) string directory to use as the model directory.
      max_checkpoints_to_keep: an (optional) integer count of saved checkpoints.
      iterations_per_loop: an (optional) integer count of log_step_count_steps.
      params: an (optional) dictionary of parameters to pass to the Estimator
        constructor.
      unused_device_fn: a device_fn to pass to RunConfig, if not use_tpu.
      master: a string necessary for TPU, pass FLAGS.master through.
      use_tpu: boolean.  set self.use_tpu if not None.
      start_from_checkpoint: string. If not None, initialize model from this
        path. According to the current implementation of Estimator, this will
        only be used in training. The inference checkpoint is loaded in a
        different place.
      session_config: a tf.ConfigProto to pass to RunConfig, if not use_tpu.
      include_debug_info: from call_variants. If True, PREDICT mode will
        include extra info such as logits and prelogits.

    Returns:
      an object implementing the tf.estimator.Estimator interface (will be a
      TPUEstimator if self.use_tpu is True).
    """
    if use_tpu is not None:
      self.use_tpu = use_tpu

    self.include_debug_info = include_debug_info

    # Set the model dir of this class to the model_dir passed in here. It's not
    # so clean but it appears to be necessary due to the way estimators are
    # constructed (i.e., model_dir is set late).
    self.model_dir = model_dir

    # These flags are exclusive if not None, and 0 means disable.
    save_checkpoints_secs = None
    save_checkpoints_steps = None
    if FLAGS.save_interval_secs >= 0:
      save_checkpoints_secs = FLAGS.save_interval_secs
    if FLAGS.save_interval_steps >= 0:
      save_checkpoints_steps = FLAGS.save_interval_steps

    params = params if params is not None else {}
    warm_start_from = self._create_warm_start_settings(start_from_checkpoint)
    if self.use_tpu:
      tpu_cfg=tpu_config.TPUConfig(
          iterations_per_loop=iterations_per_loop)
      config = tpu_config.RunConfig(
          master=master,
          evaluation_master=master,
          model_dir=model_dir,
          log_step_count_steps=iterations_per_loop,
          keep_checkpoint_max=max_checkpoints_to_keep,
          save_checkpoints_secs=save_checkpoints_secs,
          save_checkpoints_steps=save_checkpoints_steps,
          save_summary_steps=FLAGS.save_summary_steps,
          tpu_config=tpu_cfg)

      classifier = tpu_estimator.TPUEstimator(
          use_tpu=self.use_tpu,
          model_fn=self.model_fn,
          config=config,
          # TODO: enable setting these independently.
          train_batch_size=batch_size,
          eval_batch_size=batch_size,
          predict_batch_size=batch_size,
          params=params,
          warm_start_from=warm_start_from,
      )
    else:
      config = tf_estimator.RunConfig(
          model_dir=model_dir,
          log_step_count_steps=iterations_per_loop,
          keep_checkpoint_max=max_checkpoints_to_keep,
          # device_fn=device_fn,  # Not in tf1.8?
          save_checkpoints_secs=save_checkpoints_secs,
          save_checkpoints_steps=save_checkpoints_steps,
          save_summary_steps=FLAGS.save_summary_steps,
          session_config=session_config,
      )
      # The TPUEstimator interface implicitly adds batch_size to the params
      # dict. Do so explicitly here, so that we can use the same model_fn.
      params_with_batch_size = {'batch_size': batch_size}
      params_with_batch_size.update(params)

      classifier = tf_estimator.Estimator(
          model_fn=self.model_fn,
          config=config,
          params=params_with_batch_size,
          warm_start_from=warm_start_from)
    return classifier

  def model_fn(self, features, labels, mode, params):
    """A model_fn satisfying the Estimator API.

    Args:
      features: a dictionary supplying features.
      labels: a tensor of labels.
      mode: one of tf.estimator.ModeKeys.{EVAL,TRAIN}
      params: a dictionary of parameters.

    Returns:
      a tf.estimator.EstimatorSpec or tpu_estimator.TPUEstimatorSpec,
      depending on self.use_tpu.
    """
    raise NotImplementedError

  def session_eval_hooks(self):
    """Returns a list of tf.train.SessionRunHook classes.

    A typical use case is to provide a hook to load the EMA variables.

    These will be instantiated and invoked by
      eval_hooks = [
          h(model_dir) for h in model.session_eval_hooks()
      ]
      estimator.evaluate(hooks=...).

    Note that this is done according to the instructions in
    cloud_tpu/models/inception/inception_v3.py. A newer idea is in
    tpuestimator-scaffold, but we haven't tried that approach.
    """
    return []

  def session_predict_hooks(self):
    """Returns a list of tf.train.SessionRunHook classes.

    A typical use case is to provide a hook to load the EMA variables.

    These will be instantiated and invoked by
      predict_hooks = [
          h(checkpoint_path) for h in model.session_predict_hooks()
      ]
      estimator.predict(hooks=...).

    Note that this is done according to the instructions in
    cloud_tpu/models/inception/inception_v3.py. A newer idea is in
    tpuestimator-scaffold, but we haven't tried that approach.
    """
    return []

  def create(self, images, num_classes, is_training):
    """Creates a new model.

    Args:
      images: A 4-D tensor of (batch_size, height, width, channels) of pileup
        images.
      num_classes: integer. How many prediction classes are we expecting in
        model?
      is_training: boolean. Should we setup model for training (True) or for
        inference (False).

    Returns:
      A dictionary, containing string keys mapped to endpoint tensors of this
      model. The dictionary must contain a key 'Predictions' that contains the
      probability of having each of 'num_classes' classes.
    """
    try:
      return self._create(images, num_classes, is_training)
    except (ValueError, tf.errors.OpError) as e:
      if self._is_bad_image_dimension_exception(e):
        _, height, width, _ = images.get_shape().as_list()
        message = (
            'Unsupported image dimensions detected: model {} was given images '
            'of w={} x h={} but a TensorFlow exception occurred while building '
            'the model, which typically indicates those dimensions are not '
            'supported by the model. The supported dimensions for {} are {}'
        ).format(self.name, width, height, self.name,
                 self.supported_dimensions_message)
        raise UnsupportedImageDimensionsError(message)
      else:
        raise

  def _is_bad_image_dimension_exception(self, exception):
    return any(
        x in str(exception) for x in ['Negative dimension', 'SpatialSqueeze'])

  def _create(self, images, num_classes, is_training):
    """To be overloaded by subclasses to actually create the model."""
    raise NotImplementedError

  def preprocess_images(self, images):
    """Preprocessing steps needed for this model to process a batch of images.

    Args:
      images: A (batch_size, height, width, channels) 4-D Tensor of type uint8.

    Returns:
      A new batch of images, potentially with different dimensions, based on the
      input but transformed as necessary to use with this model.
    """
    raise NotImplementedError

  @property
  def is_trainable(self):
    """Returns True if this model can be trained."""
    return True

  # TODO: Add export to save representation suitable for inference.

  def __str__(self):
    return 'DeepVariantModel(name={})'.format(self.name)

  def variables_to_restore_from_model(self, exclude_scopes=None):
    """Gets the list of model variables that should be restored.

    The primary use of this function is to get a subset of tf.Variables from a
    slim-defined model that we'd like to restore from a checkpoint. The
    checkpoint generally contains all of the variables in the graph during
    training, including things like the backprop variables, moving averages for
    visualization, etc. Simply restoring all of those variables is brittle, as
    we often want to start a new training run, maybe using a different
    optimizer, different visualization variables, or replacing part of the model
    with a new classification layer, as unneeded variables from the checkpoint
    get loaded into the graph and/or new TF variables not present in the graph
    cannot be found, raising exceptions. This function allows a clean API to get
    just the *model* variables from a graph, excluding all of those non-model
    variables, along with optionally removing parts of the model graph via
    exclude scopes.

    This function calls slim.get_model_variables() to get the raw list of all
    variables associated with the MODEL_VARIABLES collection. It then filters
    away all variables that match any of the scopes in exclude_scopes. For
    example, suppose we have a model with three variables with names:

      w1 = model/l1/weight1
      w2 = model/l2/weight2
      w3 = model/l2/weight3

    Without any exclude scopes, we would return these three variables [w1, w2,
    and w3]. Providing exclude_scopes=['model/l2'] would return only [w1], while
    exclude_scopes=['model/l1'] would return [w2, w3].

    Args:
      exclude_scopes: None, or a list of strings. Each string is a scope
        specification, such as "model/l1" to match all variables whose name
        starts with "model/l1".

    Returns:
      A list of tf.Variable objects.
    """
    vars_to_include = slim.get_model_variables()
    # We aren't excluding any variables, so just return vars_to_include.
    if not exclude_scopes:
      return vars_to_include

    vars_to_exclude = set()
    for scope in exclude_scopes:
      vars_to_exclude |= set(slim.get_variables(scope))
    return [v for v in vars_to_include if v not in vars_to_exclude]


class DeepVariantSlimModel(DeepVariantModel):
  """Baseclass for DeepVariant models based on Slim networks."""

  def __init__(self, name, pretrained_model_path, n_classes_model_variable,
               excluded_scopes_for_incompatible_classes,
               excluded_scopes_for_incompatible_channels):
    """Creates an DeepVariant CNN network based on a tf.slim model.

    Args:
      name: see baseclass.
      pretrained_model_path: see baseclass.
      n_classes_model_variable: str. A fully-qualitified TF variable name in the
        model that we can use to determine the shape of the output
        classification layer of the model. For example, in inception-v3 from
        slim this is 'InceptionV3/Logits/Conv2d_1c_1x1/weights'.
      excluded_scopes_for_incompatible_classes: set of str. A set of scopes
        that will be excluded when restoring from a checkpoint to avoid loading
        incompatible #classes.
      excluded_scopes_for_incompatible_channels: set of str. A set of scopes
        that will be excluded when restoring from a checkpoint to avoid loading
        incompatible #channels.
    Raises:
      ValueError: If any of the arguments are invalid.
    """
    super(DeepVariantSlimModel, self).__init__(
        name=name, pretrained_model_path=pretrained_model_path)
    self.n_classes_model_variable = n_classes_model_variable
    self.excluded_scopes_for_incompatible_classes = (
        excluded_scopes_for_incompatible_classes)
    self.excluded_scopes_for_incompatible_channels = (
        excluded_scopes_for_incompatible_channels)

  def preprocess_images(self, images):
    """Applies preprocessing operations for Inception images.

    Because this will run in model_fn, on the accelerator, we use operations
    that efficiently execute there.

    Args:
      images: An Tensor of shape [batch_size height, width, channel] with uint8
        values.

    Returns:
      A tensor of images of shape [batch_size height, width, channel]
      containing floating point values, with all points rescaled between
      -1 and 1 and possibly resized.
    """
    images = tf.cast(images, dtype=tf.float32)
    images = tf.subtract(images, 128.0)
    images = tf.math.divide(images, 128.0)
    return images

  def model_fn(self, features, labels, mode, params):
    """A model_fn for slim (really inception_v3), satisfying the Estimator API.

    Args:
      features: a single Tensor or dict of same (from input_fn).
      labels: a single Tensor or dict of same (from input_fn).
      mode: tf.estimator.ModeKeys.
      params: dict.

    Returns:
      EstimatorSpec or TPUEstimatorSpec depending on self.use_tpu.
    """
    # NB. The basic structure of this started from
    # //third_party/cloud_tpu/models/inception/inception_v3.py

    # TODO: get this from the model.
    num_classes = dv_constants.NUM_CLASSES

    images = features['image']
    images = self.preprocess_images(images)

    endpoints = self.create(
        images=images,
        num_classes=num_classes,
        is_training=mode == tf_estimator.ModeKeys.TRAIN)

    logits = endpoints['Logits']

    predictions = endpoints
    predictions.update({
        'classes': tf.argmax(input=logits, axis=1, output_type=tf.int32),
        'probabilities': tf.nn.softmax(logits, name='softmax_tensor')
    })
    prelogits = endpoints['PreLogits'] if self.include_debug_info else None
    if mode == tf_estimator.ModeKeys.PREDICT:
      return self._model_fn_predict(mode, features, logits, prelogits=prelogits)

    # Compute loss.
    one_hot_labels = tf.one_hot(labels, num_classes, dtype=tf.int32)
    tf.compat.v1.losses.softmax_cross_entropy(
        onehot_labels=one_hot_labels,
        logits=logits,
        weights=1.0,
        label_smoothing=FLAGS.label_smoothing)
    total_loss = tf.compat.v1.losses.get_total_loss(
        add_regularization_losses=True)
    return self.make_ops_and_estimator(features, endpoints, labels, logits,
                                       predictions, total_loss, mode, params)

  def make_ops_and_estimator(self, features, endpoints, labels, logits,
                             predictions, total_loss, mode, params):
    """Make EstimatorSpec for the current model.

    Args:
      features: a single Tensor or dict of same (from input_fn).
      endpoints:  a dictionary, containing string keys mapped to endpoint
        tensors of this model. The dictionary must contain a key 'Predictions'
        that contains the probability of having each of 'num_classes' classes.
      labels: a single Tensor or dict of same (from input_fn).
      logits: a single Tensor with logits
      predictions: A dictionaty that must contain the following keys: 'Logits'
        and 'Predictions'.
      total_loss:  a single Tensor with a loss
      mode: tf.estimator.ModeKeys.
      params: dict.

    Returns:
      EstimatorSpec or TPUEstimatorSpec depending on self.use_tpu.
    """
    # Note, below, one of train_op or eval_metrics will be None, and the other
    # will be populated, depending on mode.

    # There are a lot of arguments here; that's to avoid referencing flags in
    # leaf functions.
    train_op, host_call = self._model_fn_train(
        mode=mode,
        total_loss=total_loss,
        # get() here to be robust when we are in eval mode and batches_per_epoch
        # hasn't been provided. In eval mode, model_fn_train will return without
        # doing anything.
        batches_per_epoch=params.get('batches_per_epoch', None),
        num_epochs_per_decay=FLAGS.num_epochs_per_decay,
        initial_learning_rate=FLAGS.learning_rate,
        learning_rate_decay_factor=FLAGS.learning_rate_decay_factor,
        rmsprop_decay=FLAGS.rmsprop_decay,
        rmsprop_momentum=FLAGS.rmsprop_momentum,
        rmsprop_epsilon=FLAGS.rmsprop_epsilon,
        moving_average_decay=FLAGS.moving_average_decay)

    eval_metrics = self._model_fn_eval(
        mode=mode,
        features=features,
        labels=labels,
        endpoints=endpoints,
        logits=logits,
        use_logits=False)

    spec = tpu_estimator.TPUEstimatorSpec(
        mode=mode,
        loss=total_loss,
        train_op=train_op,
        host_call=host_call,
        eval_metrics=eval_metrics,
        predictions=predictions)
    if self.use_tpu:
      return spec
    else:
      return spec.as_estimator_spec()

  def _model_fn_predict(self, mode, features, logits, prelogits=None):
    """This is the PREDICT part of model_fn."""
    assert mode == tf_estimator.ModeKeys.PREDICT
    predictions = {
        # We don't actually use classes downstream right now.
        # 'classes': tf.argmax(input=logits, axis=1, output_type=tf.int32),
        'probabilities': tf.nn.softmax(logits, name='softmax_tensor'),
        # DV2 call_variants wants these passed through.
        'variant': features['variant'],
        'alt_allele_indices': features['alt_allele_indices'],
    }
    if self.include_debug_info:
      if logits is not None:
        predictions.update({'logits': logits})
      if prelogits is not None:
        predictions.update({'prelogits': prelogits})
    if 'label' in features:
      predictions['label'] = features['label']
    if 'locus' in features:
      predictions['locus'] = features['locus']
    if self.use_tpu:
      return tpu_estimator.TPUEstimatorSpec(mode=mode, predictions=predictions)
    else:
      return tf_estimator.EstimatorSpec(mode=mode, predictions=predictions)

  def _model_fn_eval(self, mode, features, labels, endpoints, logits,
                     use_logits):
    """This is the EVAL part of model_fn."""
    if mode != tf_estimator.ModeKeys.EVAL:
      return None
    if use_logits:
      eval_predictions = logits
    else:
      eval_predictions = endpoints['Predictions']
    variant_type = features['variant_type']
    eval_metrics = (eval_metric_fn, [labels, eval_predictions, variant_type])
    if not self.use_tpu:
      for name, value in eval_metrics[0](*eval_metrics[1]).items():
        tf.compat.v1.summary.scalar(tensor=value, name=name)
    return eval_metrics

  def _model_fn_train(self, mode, total_loss, batches_per_epoch,
                      num_epochs_per_decay, initial_learning_rate,
                      learning_rate_decay_factor, rmsprop_decay,
                      rmsprop_momentum, rmsprop_epsilon, moving_average_decay):
    """This is the TRAIN part of model_fn."""
    if mode != tf_estimator.ModeKeys.TRAIN:
      return None, None

    # Configure the learning rate using an exponetial decay.
    global_step = tf.compat.v1.train.get_or_create_global_step()
    current_epoch = tf.cast(global_step, tf.float32) / batches_per_epoch
    decay_steps = int(1.0 * batches_per_epoch * num_epochs_per_decay)

    learning_rate = tf.compat.v1.train.exponential_decay(
        learning_rate=initial_learning_rate,
        global_step=global_step,
        decay_steps=decay_steps,
        decay_rate=learning_rate_decay_factor,
        staircase=True)

    # Set a minimum boundary for the learning rate to be a fixed value of 1e-9.
    # It's common to see these tf.max(...) operations when training inception,
    # with a max of 1e-4 * initial_learning_rate but this makes it hard to
    # explore learning rate schedules that decay quickly or by a lot of each
    # step. Here we just use a very small constant 1e-9 as the minimum value.
    learning_rate = tf.maximum(learning_rate, 1e-9, name='learning_rate')

    optimizer = tf.compat.v1.train.RMSPropOptimizer(
        learning_rate,
        rmsprop_decay,
        momentum=rmsprop_momentum,
        epsilon=rmsprop_epsilon)
    if self.use_tpu:
      optimizer = tpu_optimizer.CrossShardOptimizer(optimizer)
    update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
      train_op = optimizer.minimize(total_loss, global_step=global_step)

    # NB. In the inception code this was "tf.trainable_variables()
    # + tf.moving_average_variables()", but we've settled on just
    # tf.model_variables() in the existing production DV2.
    variables_to_average = tf.compat.v1.model_variables()
    variable_averages = tf.train.ExponentialMovingAverage(
        decay=moving_average_decay, num_updates=global_step)
    with tf.control_dependencies([train_op
                                 ]), tf.compat.v1.name_scope('moving_average'):
      train_op = variable_averages.apply(variables_to_average)
    tf.compat.v1.add_to_collection(tf.compat.v1.GraphKeys.UPDATE_OPS, train_op)

    # Compute the current epoch and associated learning rate from global_step.
    metric_dict = {
        'current_epoch': current_epoch,
        'total_loss': total_loss,
        'learning_rate': learning_rate,
    }
    host_call = self.construct_scalar_host_call(
        metric_dict=metric_dict, model_dir=self.model_dir, prefix='training/')

    return train_op, host_call

  def session_eval_hooks(self):
    return [LoadEMAHook]

  def session_predict_hooks(self):
    return [PredictEMAHook]


class DeepVariantInceptionV3(DeepVariantSlimModel):
  """DeepVariant inception_v3 network."""

  def __init__(self):
    """Creates an inception-v3 network for DeepVariant."""
    super(DeepVariantInceptionV3, self).__init__(
        name='inception_v3',
        n_classes_model_variable='InceptionV3/Logits/Conv2d_1c_1x1/weights',
        excluded_scopes_for_incompatible_classes=[
            'InceptionV3/Logits', 'InceptionV3/Conv2d_1a_3x3'
        ],
        excluded_scopes_for_incompatible_channels=['InceptionV3/Conv2d_1a_3x3'],
        pretrained_model_path=('/namespace/vale-project/models/classification/'
                               'imagenet/inception_v3/model.ckpt-9591376'))
    self.supported_dimensions_message = (
        'odd widths between 75-361 and any heights between 75-362')

  def _create(self, images, num_classes, is_training):
    """See baseclass."""
    with slim.arg_scope(inception_v3.inception_v3_arg_scope()):
      _, endpoints = inception_v3.inception_v3(
          images, num_classes, create_aux_logits=False, is_training=is_training)
      return endpoints


class DeepVariantInceptionV3Embedding(DeepVariantInceptionV3):
  """DeepVariant inception_v3_embedding network."""

  def __init__(self):
    """Creates an inception_v3_embedding network for DeepVariant."""
    super(DeepVariantInceptionV3Embedding, self).__init__()
    self.name = 'inception_v3_embedding'
    # vocab_size should be a number larger than the number of sequencing types
    self.vocab_size = 5
    self.embedding_size = 200
    self.dropout_keep_prob = 0.8

  def _create(self, inputs, num_classes, is_training):
    """Creates a new inception_v3_embedding model.

    Args:
      inputs: A tuple of two elements (images, sequencing_types). images is a
        4-D tensor of (batch_size, height, width, channels) of pileup images.
        sequencing_types is a 1-D tensor of (batch_size) of example sequencing
        types.
      num_classes: integer. How many prediction classes are we expecting in
        model?
      is_training: boolean. Should we setup model for training (True) or for
        inference (False).

    Returns:
      A dictionary, containing string keys mapped to endpoint tensors of this
      model.
    """
    images, sequencing_type = inputs
    endpoints = super(DeepVariantInceptionV3Embedding,
                      self)._create(images, num_classes, is_training)

    with tf.compat.v1.variable_scope('Embeddings'):
      # Take the graph all the way till PreLogits
      net = endpoints['PreLogits']
      net = slim.flatten(net)
      embeddings = self._create_embeddings(sequencing_type)
      net = tf.concat([net, embeddings], 1)

      endpoints['Embeddings'] = net

    with tf.compat.v1.variable_scope('Logits'):
      if isinstance(net.shape[1], int):
        hidden_size = net.shape[1] // 2
      else:
        hidden_size = net.shape[1].value // 2

      net = slim.fully_connected(net, hidden_size, activation_fn=None)
      # TODO: Explore using ReLU before norm
      net = slim.layer_norm(net, scale=False, activation_fn=tf.nn.relu)
      net = slim.dropout(net, self.dropout_keep_prob, is_training=is_training)
      net = slim.fully_connected(net, num_classes, activation_fn=None)

      endpoints.update({'Logits': net, 'Predictions': tf.nn.softmax(net)})

    return endpoints

  def _create_embeddings(self, indices):
    """Create word embeddings."""
    embeddings = self._embedding_lookup(indices)
    embeddings = slim.fully_connected(
        embeddings, self.embedding_size, activation_fn=None)
    return embeddings

  def _embedding_lookup(self, input_ids, word_embedding_name='seq_type_emb'):
    """Looks up words embeddings for id tensor.

    Args:
      input_ids: int64 Tensor of shape [batch_size, ] containing word ids.
      word_embedding_name: string. Name of the embedding table.

    Returns:
      float Tensor of shape [batch_size, embedding_size].
    """
    embedding_table = tf.compat.v1.get_variable(
        name=word_embedding_name,
        shape=[self.vocab_size, self.embedding_size],
        initializer=tf.compat.v1.keras.initializers.VarianceScaling(
            scale=1.0, mode='fan_avg', distribution='uniform'),
        collections=[
            tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES,
            tf.compat.v1.GraphKeys.MODEL_VARIABLES,
            tf.compat.v1.GraphKeys.GLOBAL_VARIABLES
        ])

    return tf.nn.embedding_lookup(params=embedding_table, ids=input_ids)

  def model_fn(self, features, labels, mode, params):
    """A model_fn for slim, satisfying the Estimator API.

    Args:
      features: a single Tensor or dict of same (from input_fn).
      labels: a single Tensor or dict of same (from input_fn).
      mode: tf.estimator.ModeKeys.
      params: dict.

    Returns:
      EstimatorSpec or TPUEstimatorSpec depending on self.use_tpu.

    Raises:
      ValueError: if FLAGS.seq_type_embedding_size is not positive.
    """
    # NB. The basic structure of this started from
    # //third_party/cloud_tpu/models/inception/inception_v3.py

    # TODO: get this from the model.
    num_classes = dv_constants.NUM_CLASSES

    if FLAGS.seq_type_embedding_size <= 0:
      raise ValueError(
          'Expected seq_type_embedding_size to be a positive number but saw %i '
          'instead.' % FLAGS.seq_type_embedding_size)
    self.embedding_size = FLAGS.seq_type_embedding_size

    images = features['image']
    images = self.preprocess_images(images)
    sequencing_type = features['sequencing_type']

    endpoints = self.create(
        images=(images, sequencing_type),
        num_classes=num_classes,
        is_training=mode == tf_estimator.ModeKeys.TRAIN)

    logits = endpoints['Logits']

    predictions = endpoints
    predictions.update({
        'classes': tf.argmax(input=logits, axis=1, output_type=tf.int32),
        'probabilities': tf.nn.softmax(logits, name='softmax_tensor')
    })

    prelogits = endpoints['PreLogits'] if self.include_debug_info else None
    if mode == tf_estimator.ModeKeys.PREDICT:
      return self._model_fn_predict(mode, features, logits, prelogits=prelogits)

    # Compute loss.
    one_hot_labels = tf.one_hot(labels, num_classes, dtype=tf.int32)
    tf.compat.v1.losses.softmax_cross_entropy(
        onehot_labels=one_hot_labels,
        logits=logits,
        weights=1.0,
        label_smoothing=FLAGS.label_smoothing)
    total_loss = tf.compat.v1.losses.get_total_loss(
        add_regularization_losses=True)

    return self.make_ops_and_estimator(features, endpoints, labels, logits,
                                       predictions, total_loss, mode, params)


class DeepVariantAttentionInceptionV3(DeepVariantSlimModel):
  """DeepVariant inception_v3 network."""

  def __init__(self, attention_module='', attention_position='all'):
    """Creates an inception-v3 network for DeepVariant."""
    super(DeepVariantAttentionInceptionV3, self).__init__(
        name='attention_inception_v3',
        n_classes_model_variable='InceptionV3/Logits/Conv2d_1c_1x1/weights',
        excluded_scopes_for_incompatible_classes=[
            'InceptionV3/Logits', 'InceptionV3/Conv2d_1a_3x3'
        ],
        excluded_scopes_for_incompatible_channels=['InceptionV3/Conv2d_1a_3x3'],
        pretrained_model_path=('/namespace/vale-project/models/classification/'
                               'imagenet/inception_v3/model.ckpt-9591376'))
    self.supported_dimensions_message = (
        'odd widths between 75-361 and any heights between 75-362')

    self.attention_module = attention_module
    self.attention_position = attention_position

  def _create(self, images, num_classes, is_training):
    """See baseclass."""
    with slim.arg_scope(
        attention_inception_v3.attention_inception_v3_arg_scope()):
      _, endpoints = attention_inception_v3.attention_inception_v3(
          images,
          num_classes,
          create_aux_logits=False,
          is_training=is_training,
          attention_module=self.attention_module,
          attention_position=self.attention_position)
    return endpoints


class DeepVariantPlaceholderModel(DeepVariantModel):
  """BaseClass for placeholder models that are useful for testing and benchmarking."""

  def __init__(self, name):
    """Creates a Placeholder model."""
    # Note the pretrained model path isn't used but we must return a valid
    # string so here we just return "UNUSED".
    super(DeepVariantPlaceholderModel, self).__init__(
        name=name, pretrained_model_path='UNUSED')

  def preprocess_images(self, images):
    """Preprocess images for placeholder model."""
    # Note these calculations aren't necessary, but they are included here to
    # mimic the data processing pipeline used by inception. We may consider
    # removing them in a future CL, or making them optional, to reduce CPU cost
    # of this model.
    images = tf.cast(images, dtype=tf.float32)
    images = tf.subtract(images, 128.0)
    images = tf.math.divide(images, 128.0)
    return images

  @property
  def is_trainable(self):
    """A placeholder model cannot be trained."""
    return False


class DeepVariantConstantModel(DeepVariantPlaceholderModel):
  """Returns a constant probability distribution for each example."""

  def __init__(self, predictions=None):
    """Creates a constant model.

    Args:
      predictions: list[float]. Values to return for Predictions, which should
        be a floatting point value between 0 and 1 for each class, normalized so
        the sum of the values is 1. Predictions should have dimension
        [num_classes].

    Raises:
      ValueError: if sum(predictions) is not close to 1.
    """
    # Note the pretrained model path isn't used but we must return a valid
    # string so here we just return "UNUSED".
    super(DeepVariantConstantModel, self).__init__(name='constant')
    if predictions is None:
      self.predictions = [0.0, 1.0, 0.0]
    elif math.abs(sum(predictions) - 1) > 1e-6:
      raise ValueError('Sum of predictions should be ~1', predictions)
    else:
      self.predictions = predictions

  @staticmethod
  def _predictions(pred_const, batch_size):
    return {
        'Predictions':
            tf.reshape(
                tf.tile(pred_const, [batch_size]),
                shape=(batch_size, tf.shape(input=pred_const)[0]))
    }

  def _create(self, images, num_classes, is_training):
    assert num_classes == len(self.predictions)
    batch_size = tf.shape(input=images)[0]
    pred_const = tf.constant(self.predictions)
    return self._predictions(pred_const, batch_size)

  def model_fn(self, features, labels, mode, params):
    """A model_fn for the constant model."""
    if mode == tf_estimator.ModeKeys.PREDICT:
      batch_size = tf.shape(input=features['image'])[0]
      logging.info('actual_batch_size %s', batch_size)
    else:
      batch_size = params['batch_size']
      logging.info('batch_size %s', batch_size)
    pred_const = tf.constant(self.predictions)
    endpoints = self._predictions(pred_const, batch_size)
    encoded_variants = features['variant']
    # For the constant model, which is for testing only, we just set the
    # variant_types to 0s. This is needed because it doesn't work to fetch
    # 'variant_type' from either features or endpoints here. Annoying.
    # variant_types = features['variant_type']     # Fails.
    # variant_types = endpoints['variant_type']    # Fails.
    variant_types = tf.zeros(shape=(batch_size,), dtype=tf.int64)

    if mode == tf_estimator.ModeKeys.PREDICT:
      predictions = {
          'probabilities': endpoints['Predictions'],
          'variant': encoded_variants,
          'alt_allele_indices': features['alt_allele_indices'],
      }
      endpoints.update(predictions)

    if mode == tf_estimator.ModeKeys.EVAL:
      eval_metrics = (eval_metric_fn,
                      [labels, endpoints['Predictions'], variant_types])
    else:
      eval_metrics = None

    loss = tf.constant(0.0)
    train_op = None

    spec = tpu_estimator.TPUEstimatorSpec(
        mode=mode,
        loss=loss,
        train_op=train_op,
        eval_metrics=eval_metrics,
        predictions=endpoints)
    if self.use_tpu:
      return spec
    else:
      return spec.as_estimator_spec()


class DeepVariantSmallModel(DeepVariantSlimModel):
  """A smaller of version of the DeepVariant model.

     Uses only the first layers of Inception net.
  """

  def __init__(self, representation_layer='Mixed_5d'):
    """Creates an DeepVariant CNN network based on a tf.slim model.

    Args:
      representation_layer: string. The name of the layer from the Inception net
        which will be used as an endpoint.

    Raises:
      ValueError: If any of the arguments are invalid.
    """
    super(DeepVariantSmallModel, self).__init__(
        name='small_inception',
        pretrained_model_path=('/namespace/vale-project/models/classification/'
                               'imagenet/inception_v3/model.ckpt-9591376'),
        n_classes_model_variable='InceptionV3/Logits/Conv2d_1c_1x1/weights',
        excluded_scopes_for_incompatible_classes=[
            'InceptionV3/Logits', 'InceptionV3/Conv2d_1a_3x3'
        ],
        excluded_scopes_for_incompatible_channels=['InceptionV3/Conv2d_1a_3x3'])

    self.representation_layer = representation_layer

  def model_fn(self, features, labels, mode, params):
    """A model_fn for slim (really inception_v3), satisfying the Estimator API.

    Args:
      features: a single Tensor or dict of same (from input_fn).
      labels: a single Tensor or dict of same (from input_fn).
      mode: tf.estimator.ModeKeys.
      params: dict.

    Returns:
      EstimatorSpec or TPUEstimatorSpec depending on self.use_tpu.

    Raises:
      ValueError: If representation_layer was not found in the Inception
      architecture
    """
    # NB. The basic structure of this started from
    # //third_party/cloud_tpu/models/inception/inception_v3.py

    # TODO: get this from the model.
    num_classes = dv_constants.NUM_CLASSES

    images = features['image']
    images = self.preprocess_images(images)

    endpoints = self.create(
        images=images,
        num_classes=num_classes,
        is_training=mode == tf_estimator.ModeKeys.TRAIN)

    if self.representation_layer not in endpoints.keys():
      raise ValueError('Layer {} is not found Inception endpoints.'
                       'Available Inception net endpoints: {}'.format(
                           self.representation_layer, endpoints.keys()))

    mid_layer = endpoints[self.representation_layer]
    # Perform 1x1 convolution similarly to the Inception architecture
    # (see 'Predictions' end points in inception_v3 architecture)

    tower = tf.nn.conv2d(
        mid_layer, 1, [1, 1], stride=1, activation_fn=tf.nn.relu)

    batch_size = tower.get_shape()[0].value
    tower = tf.reshape(tower, [batch_size, -1])

    with tf.compat.v1.variable_scope('denselayers'):
      with slim.arg_scope([slim.fully_connected], activation_fn=tf.nn.relu):
        logits = slim.fully_connected(tower, num_classes, scope='Dense')

    predictions = endpoints
    predictions.update({
        'classes': tf.argmax(input=logits, axis=1, output_type=tf.int32),
        'probabilities': tf.nn.softmax(logits, name='softmax_tensor'),
        'Logits': logits,
        'Predictions': slim.softmax(logits)
    })

    if mode == tf_estimator.ModeKeys.PREDICT:
      return self._model_fn_predict(mode, features, logits)

    # Compute loss.
    one_hot_labels = tf.one_hot(labels, num_classes, dtype=tf.int32)
    tf.compat.v1.losses.softmax_cross_entropy(
        onehot_labels=one_hot_labels,
        logits=logits,
        weights=1.0,
        label_smoothing=FLAGS.label_smoothing)
    total_loss = tf.compat.v1.losses.get_total_loss(
        add_regularization_losses=True)

    return self.make_ops_and_estimator(features, endpoints, labels, logits,
                                       predictions, total_loss, mode, params)

  def _create(self, images, num_classes, is_training):
    """See baseclass."""
    with slim.arg_scope(inception_v3.inception_v3_arg_scope()):
      _, endpoints = inception_v3.inception_v3(
          images, num_classes, create_aux_logits=False, is_training=is_training)
      return endpoints


# Our list of pre-defined model classes.
_MODEL_CLASSES = [
    DeepVariantSmallModel, DeepVariantInceptionV3, DeepVariantConstantModel,
    DeepVariantInceptionV3Embedding, DeepVariantAttentionInceptionV3
]


def all_models():
  """Gets a list of the all of the known model classes."""
  return list(_MODEL_CLASSES)


def production_models():
  """Gets a list of the models that we test extensively."""
  return [get_model('inception_v3'), get_model('inception_v3_embedding')]


def get_model(model_name, **kwargs):
  """Looks up a DeepVariantModel by name.

  Args:
    model_name: String. Looks for a pre-defined DeepVariantModel with a name
      equal to this model_name string.
    **kwargs: arguments to pass to model constructor.

  Returns:
    A DeepVariantModel instance.

  Raises:
    ValueError: If no model exists with model_name.
  """
  for model_class in _MODEL_CLASSES:
    # Instantiate the model with any provided arguments.
    if kwargs:
      model = DeepVariantAttentionInceptionV3(**kwargs)
    else:
      model = model_class()
    if model_name == model.name:
      return model
  raise ValueError('Unknown model_name {}, options are {}'.format(
      model_name, [model_class().name for model_class in _MODEL_CLASSES]))