model.py

# Copyright 2019 DeepMind Technologies Limited
#
# 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
#
#     https://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.

"""Open Source Version of the Hierarchical Probabilistic U-Net."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import geco_utils
import sonnet as snt
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import unet_utils


class _HierarchicalCore(snt.AbstractModule):
  """A U-Net encoder-decoder with a full encoder and a truncated decoder.

  The truncated decoder is interleaved with the hierarchical latent space and
  has as many levels as there are levels in the hierarchy plus one additional
  level.
  """

  def __init__(self, latent_dims, channels_per_block,
               down_channels_per_block=None, activation_fn=tf.nn.relu,
               initializers=None, regularizers=None, convs_per_block=3,
               blocks_per_level=3, name='HierarchicalDecoderDist'):
    """Initializes a HierarchicalCore.

    Args:
      latent_dims: List of integers specifying the dimensions of the latents at
        each scale. The length of the list indicates the number of U-Net decoder
        scales that have latents.
      channels_per_block: A list of integers specifying the number of output
        channels for each encoder block.
      down_channels_per_block: A list of integers specifying the number of
        intermediate channels for each encoder block or None. If None, the
        intermediate channels are chosen equal to channels_per_block.
      activation_fn: A callable activation function.
      initializers: Optional dict containing ops to initialize the filters (with
        key 'w') or biases (with key 'b'). The default initializer for the
        weights is a truncated normal initializer, which is commonly used when
        the inputs are zero centered (see
        https://arxiv.org/pdf/1502.03167v3.pdf). The default initializer for the
        bias is a zero initializer.
      regularizers: Optional dict containing regularizers for the filters
        (with key 'w') and the biases (with key 'b'). As a default, no
        regularizers are used. A regularizer should be a function that takes a
        single `Tensor` as an input and returns a scalar `Tensor` output, e.g.
        the L1 and L2 regularizers in `tf.contrib.layers`.
      convs_per_block: An integer specifying the number of convolutional layers.
      blocks_per_level: An integer specifying the number of residual blocks per
        level.
      name: A string specifying the name of the module.
    """
    super(_HierarchicalCore, self).__init__(name=name)
    self._latent_dims = latent_dims
    self._channels_per_block = channels_per_block
    self._activation_fn = activation_fn
    self._initializers = initializers
    self._regularizers = regularizers
    self._convs_per_block = convs_per_block
    self._blocks_per_level = blocks_per_level
    if down_channels_per_block is None:
      self._down_channels_per_block = channels_per_block
    else:
      self._down_channels_per_block = down_channels_per_block
    self._name = name

  def _build(self, inputs, mean=False, z_q=None):
    """A build-method allowing to sample from the module as specified.

    Args:
      inputs: A tensor of shape (b,h,w,c). When using the module as a prior the
      `inputs` tensor should be a batch of images. When using it as a posterior
      the tensor should be a (batched) concatentation of images and
      segmentations.
      mean: A boolean or a list of booleans. If a boolean, it specifies whether
        or not to use the distributions' means in ALL latent scales. If a list,
        each bool therein specifies whether or not to use the scale's mean. If
        False, the latents of the scale are sampled.
      z_q: None or a list of tensors. If not None, z_q provides external latents
        to be used instead of sampling them. This is used to employ posterior
        latents in the prior during training. Therefore, if z_q is not None, the
        value of `mean` is ignored. If z_q is None, either the distributions
        mean is used (in case `mean` for the respective scale is True) or else
        a sample from the distribution is drawn.
    Returns:
      A Dictionary holding the output feature map of the truncated U-Net
      decoder under key 'decoder_features', a list of the U-Net encoder features
      produced at the end of each encoder scale under key 'encoder_outputs', a
      list of the predicted distributions at each scale under key
      'distributions', a list of the used latents at each scale under the key
      'used_latents'.
    """
    encoder_features = inputs
    encoder_outputs = []
    num_levels = len(self._channels_per_block)
    num_latent_levels = len(self._latent_dims)
    if isinstance(mean, bool):
      mean = [mean] * num_latent_levels
    distributions = []
    used_latents = []

    # Iterate the descending levels in the U-Net encoder.
    for level in range(num_levels):
      # Iterate the residual blocks in each level.
      for _ in range(self._blocks_per_level):
        encoder_features = unet_utils.res_block(
            input_features=encoder_features,
            n_channels=self._channels_per_block[level],
            n_down_channels=self._down_channels_per_block[level],
            activation_fn=self._activation_fn,
            initializers=self._initializers,
            regularizers=self._regularizers,
            convs_per_block=self._convs_per_block)

      encoder_outputs.append(encoder_features)
      if level != num_levels - 1:
        encoder_features = unet_utils.resize_down(encoder_features, scale=2)

    # Iterate the ascending levels in the (truncated) U-Net decoder.
    decoder_features = encoder_outputs[-1]
    for level in range(num_latent_levels):

      # Predict a Gaussian distribution for each pixel in the feature map.
      latent_dim = self._latent_dims[level]
      mu_logsigma = snt.Conv2D(
          2 * latent_dim,
          (1, 1),
          padding='SAME',
          initializers=self._initializers,
          regularizers=self._regularizers,
          )(decoder_features)

      mu = mu_logsigma[Ellipsis, :latent_dim]
      logsigma = mu_logsigma[Ellipsis, latent_dim:]
      dist = tfd.MultivariateNormalDiag(loc=mu, scale_diag=tf.exp(logsigma))
      distributions.append(dist)

      # Get the latents to condition on.
      if z_q is not None:
        z = z_q[level]
      elif mean[level]:
        z = dist.loc
      else:
        z = dist.sample()
      used_latents.append(z)

      # Concat and upsample the latents with the previous features.
      decoder_output_lo = tf.concat([z, decoder_features], axis=-1)
      decoder_output_hi = unet_utils.resize_up(decoder_output_lo, scale=2)
      decoder_features = tf.concat(
          [decoder_output_hi, encoder_outputs[::-1][level + 1]], axis=-1)

      # Iterate the residual blocks in each level.
      for _ in range(self._blocks_per_level):
        decoder_features = unet_utils.res_block(
            input_features=decoder_features,
            n_channels=self._channels_per_block[::-1][level + 1],
            n_down_channels=self._down_channels_per_block[::-1][level + 1],
            activation_fn=self._activation_fn,
            initializers=self._initializers,
            regularizers=self._regularizers,
            convs_per_block=self._convs_per_block)

    return {'decoder_features': decoder_features,
            'encoder_features': encoder_outputs,
            'distributions': distributions,
            'used_latents': used_latents}


class _StitchingDecoder(snt.AbstractModule):
  """A module that completes the truncated U-Net decoder.

  Using the output of the HierarchicalCore this module fills in the missing
  decoder levels such that together the two form a symmetric U-Net.
  """

  def __init__(self, latent_dims, channels_per_block, num_classes,
               down_channels_per_block=None, activation_fn=tf.nn.relu,
               initializers=None, regularizers=None, convs_per_block=3,
               blocks_per_level=3, name='StitchingDecoder'):
    """Initializes a StichtingDecoder.

    Args:
      latent_dims: List of integers specifying the dimensions of the latents at
        each scale. The length of the list indicates the number of U-Net
        decoder scales that have latents.
      channels_per_block: A list of integers specifying the number of output
        channels for each encoder block.
      num_classes: An integer specifying the number of segmentation classes.
      down_channels_per_block: A list of integers specifying the number of
        intermediate channels for each encoder block. If None, the
        intermediate channels are chosen equal to channels_per_block.
      activation_fn: A callable activation function.
      initializers: Optional dict containing ops to initialize the filters (with
        key 'w') or biases (with key 'b'). The default initializer for the
        weights is a truncated normal initializer, which is commonly used when
        the inputs are zero centered (see
        https://arxiv.org/pdf/1502.03167v3.pdf). The default initializer for the
        bias is a zero initializer.
      regularizers: Optional dict containing regularizers for the filters
        (with key 'w') and the biases (with key 'b'). As a default, no
        regularizers are used. A regularizer should be a function that takes a
        single `Tensor` as an input and returns a scalar `Tensor` output, e.g.
        the L1 and L2 regularizers in `tf.contrib.layers`.
      convs_per_block: An integer specifying the number of convolutional layers.
      blocks_per_level: An integer specifying the number of residual blocks per
        level.
      name: A string specifying the name of the module.
    """
    super(_StitchingDecoder, self).__init__(name=name)
    self._latent_dims = latent_dims
    self._channels_per_block = channels_per_block
    self._num_classes = num_classes
    self._activation_fn = activation_fn
    self._initializers = initializers
    self._regularizers = regularizers
    self._convs_per_block = convs_per_block
    self._blocks_per_level = blocks_per_level
    if down_channels_per_block is None:
      down_channels_per_block = channels_per_block
    self._down_channels_per_block = down_channels_per_block

  def _build(self, encoder_features, decoder_features):
    """Build-method that returns the segmentation logits.

    Args:
      encoder_features: A list of tensors of shape (b,h_i,w_i,c_i).
      decoder_features: A tensor of shape (b,h,w,c).
    Returns:
      Logits, i.e. a tensor of shape (b,h,w,num_classes).
    """
    num_latents = len(self._latent_dims)
    start_level = num_latents + 1
    num_levels = len(self._channels_per_block)

    for level in range(start_level, num_levels, 1):
      decoder_features = unet_utils.resize_up(decoder_features, scale=2)
      decoder_features = tf.concat([decoder_features,
                                    encoder_features[::-1][level]], axis=-1)
      for _ in range(self._blocks_per_level):
        decoder_features = unet_utils.res_block(
            input_features=decoder_features,
            n_channels=self._channels_per_block[::-1][level],
            n_down_channels=self._down_channels_per_block[::-1][level],
            activation_fn=self._activation_fn,
            initializers=self._initializers,
            regularizers=self._regularizers,
            convs_per_block=self._convs_per_block)

    return snt.Conv2D(output_channels=self._num_classes,
                      kernel_shape=(1, 1),
                      padding='SAME',
                      initializers=self._initializers,
                      regularizers=self._regularizers,
                      name='logits')(decoder_features)


class HierarchicalProbUNet(snt.AbstractModule):
  """A Hierarchical Probabilistic U-Net."""

  def __init__(self,
               latent_dims=(1, 1, 1, 1),
               channels_per_block=None,
               num_classes=2,
               down_channels_per_block=None,
               activation_fn=tf.nn.relu,
               initializers=None,
               regularizers=None,
               convs_per_block=3,
               blocks_per_level=3,
               loss_kwargs=None,
               name='HPUNet'):
    """Initializes a HierarchicalProbUNet.

    The default values are set as for the LIDC-IDRI experiments in
    `A Hierarchical Probabilistic U-Net for Modeling Multi-Scale Ambiguities',
    see https://arxiv.org/abs/1905.13077.
    Args:
      latent_dims: List of integers specifying the dimensions of the latents at
        each scales. The length of the list indicates the number of U-Net
        decoder scales that have latents.
      channels_per_block: A list of integers specifying the number of output
        channels for each encoder block.
      num_classes: An integer specifying the number of segmentation classes.
      down_channels_per_block: A list of integers specifying the number of
        intermediate channels for each encoder block. If None, the
        intermediate channels are chosen equal to channels_per_block.
      activation_fn: A callable activation function.
      initializers: Optional dict containing ops to initialize the filters (with
        key 'w') or biases (with key 'b').
      regularizers: Optional dict containing regularizers for the filters
        (with key 'w') and the biases (with key 'b').
      convs_per_block: An integer specifying the number of convolutional layers.
      blocks_per_level: An integer specifying the number of residual blocks per
        level.
      loss_kwargs: None or dictionary specifying the loss setup.
      name: A string specifying the name of the module.
    """
    super(HierarchicalProbUNet, self).__init__(name=name)
    base_channels = 24
    default_channels_per_block = (
        base_channels, 2 * base_channels, 4 * base_channels, 8 * base_channels,
        8 * base_channels, 8 * base_channels, 8 * base_channels,
        8 * base_channels
    )
    if channels_per_block is None:
      channels_per_block = default_channels_per_block
    if down_channels_per_block is None:
      down_channels_per_block =\
        tuple([i / 2 for i in default_channels_per_block])
    if initializers is None:
      initializers = {
          'w': tf.orthogonal_initializer(gain=1.0, seed=None),
          'b': tf.truncated_normal_initializer(stddev=0.001)
      }
    if regularizers is None:
      regularizers = {
          'w': tf.keras.regularizers.l2(1e-5),
          'b': tf.keras.regularizers.l2(1e-5)
      }
    if loss_kwargs is None:
      self._loss_kwargs = {
          'type': 'geco',
          'top_k_percentage': 0.02,
          'deterministic_top_k': False,
          'kappa': 0.05,
          'decay': 0.99,
          'rate': 1e-2,
          'beta': None
      }
    else:
      self._loss_kwargs = loss_kwargs
    if down_channels_per_block is None:
      down_channels_per_block = channels_per_block

    with self._enter_variable_scope():
      self._prior = _HierarchicalCore(
          latent_dims=latent_dims,
          channels_per_block=channels_per_block,
          down_channels_per_block=down_channels_per_block,
          activation_fn=activation_fn,
          initializers=initializers,
          regularizers=regularizers,
          convs_per_block=convs_per_block,
          blocks_per_level=blocks_per_level,
          name='prior')

      self._posterior = _HierarchicalCore(
          latent_dims=latent_dims,
          channels_per_block=channels_per_block,
          down_channels_per_block=down_channels_per_block,
          activation_fn=activation_fn,
          initializers=initializers,
          regularizers=regularizers,
          convs_per_block=convs_per_block,
          blocks_per_level=blocks_per_level,
          name='posterior')

      self._f_comb = _StitchingDecoder(
          latent_dims=latent_dims,
          channels_per_block=channels_per_block,
          num_classes=num_classes,
          down_channels_per_block=down_channels_per_block,
          activation_fn=activation_fn,
          initializers=initializers,
          regularizers=regularizers,
          convs_per_block=convs_per_block,
          blocks_per_level=blocks_per_level,
          name='f_comb')

      if self._loss_kwargs['type'] == 'geco':
        self._moving_average = geco_utils.MovingAverage(
            decay=self._loss_kwargs['decay'], differentiable=True,
            name='ma_test')
        self._lagmul = geco_utils.LagrangeMultiplier(
            rate=self._loss_kwargs['rate'])
      self._cache = ()

  def _build(self, seg, img):
    """Inserts all ops used during training into the graph exactly once.

    The first time this method is called given the input pair (seg, img) all
    ops relevant for training are inserted into the graph. Calling this method
    more than once does not re-insert the modules into the graph (memoization),
    thus preventing multiple forward passes of submodules for the same inputs.
    The method is private and called when setting up the loss.
    Args:
      seg: A tensor of shape (b, h, w, num_classes).
      img: A tensor of shape (b, h, w, c)
    Returns: None
    """
    inputs = (seg, img)
    if self._cache == inputs:
      return
    else:
      self._q_sample = self._posterior(
          tf.concat([seg, img], axis=-1), mean=False)
      self._q_sample_mean = self._posterior(
          tf.concat([seg, img], axis=-1), mean=True)
      self._p_sample = self._prior(
          img, mean=False, z_q=None)
      self._p_sample_z_q = self._prior(
          img, z_q=self._q_sample['used_latents'])
      self._p_sample_z_q_mean = self._prior(
          img, z_q=self._q_sample_mean['used_latents'])
      self._cache = inputs
      return

  def sample(self, img, mean=False, z_q=None):
    """Sample a segmentation from the prior, given an input image.

    Args:
      img: A tensor of shape (b, h, w, c).
      mean: A boolean or a list of booleans. If a boolean, it specifies whether
        or not to use the distributions' means in ALL latent scales. If a list,
        each bool therein specifies whether or not to use the scale's mean. If
        False, the latents of the scale are sampled.
      z_q: None or a list of tensors. If not None, z_q provides external latents
        to be used instead of sampling them. This is used to employ posterior
        latents in the prior during training. Therefore, if z_q is not None, the
        value of `mean` is ignored. If z_q is None, either the distributions
        mean is used (in case `mean` for the respective scale is True) or else
        a sample from the distribution is drawn
    Returns:
      A segmentation tensor of shape (b, h, w, num_classes).
    """
    prior_out = self._prior(img, mean, z_q)
    encoder_features = prior_out['encoder_features']
    decoder_features = prior_out['decoder_features']
    return self._f_comb(encoder_features=encoder_features,
                        decoder_features=decoder_features)

  def reconstruct(self, seg, img, mean=False):
    """Reconstruct a segmentation using the posterior.

    Args:
      seg: A tensor of shape (b, h, w, num_classes).
      img: A tensor of shape (b, h, w, c).
      mean: A boolean, specifying whether to sample from the full hierarchy of
       the posterior or use the posterior means at each scale of the hierarchy.
    Returns:
      A segmentation tensor of shape (b,h,w,num_classes).
    """
    self._build(seg, img)
    if mean:
      prior_out = self._p_sample_z_q_mean
    else:
      prior_out = self._p_sample_z_q
    encoder_features = prior_out['encoder_features']
    decoder_features = prior_out['decoder_features']
    return self._f_comb(encoder_features=encoder_features,
                        decoder_features=decoder_features)

  def rec_loss(self, seg, img, mask=None, top_k_percentage=None,
               deterministic=True):
    """Cross-entropy reconstruction loss employed in the ELBO-/ GECO-objective.

    Args:
      seg: A tensor of shape (b, h, w, num_classes).
      img: A tensor of shape (b, h, w, c).
      mask: A mask of shape (b, h, w) or None. If None no pixels are masked in
       the loss.
      top_k_percentage: None or a float in (0.,1.]. If None, a standard
        cross-entropy loss is calculated.
      deterministic: A Boolean indicating whether or not to produce the
        prospective top-k mask deterministically.
    Returns:
      A dictionary holding the mean and the pixelwise sum of the loss for the
      batch as well as the employed loss mask.
    """
    reconstruction = self.reconstruct(seg, img, mean=False)
    return geco_utils.ce_loss(
        reconstruction, seg, mask, top_k_percentage, deterministic)

  def kl(self, seg, img):
    """Kullback-Leibler divergence between the posterior and the prior.

    Args:
      seg: A tensor of shape (b, h, w, num_classes).
      img: A tensor of shape (b, h, w, c).
    Returns:
      A dictionary with keys indexing the hierarchy's levels and corresponding
      values holding the KL-term for each level (per batch).
    """
    self._build(seg, img)
    posterior_out = self._q_sample
    prior_out = self._p_sample_z_q

    q_dists = posterior_out['distributions']
    p_dists = prior_out['distributions']

    kl = {}
    for level, (q, p) in enumerate(zip(q_dists, p_dists)):
      # Shape (b, h, w).
      kl_per_pixel = tfd.kl_divergence(q, p)
      # Shape (b,).
      kl_per_instance = tf.reduce_sum(kl_per_pixel, axis=[1, 2])
      # Shape (1,).
      kl[level] = tf.reduce_mean(kl_per_instance)
    return kl

  def loss(self, seg, img, mask):
    """The full training objective, either ELBO or GECO.

    Args:
      seg: A tensor of shape (b, h, w, num_classes).
      img: A tensor of shape (b, h, w, c).
      mask: A mask of shape (b, h, w) or None. If None no pixels are masked in
       the loss.
    Returns:
      A dictionary holding the loss (with key 'loss') and the tensorboard
      summaries (with key 'summaries').
    """
    summaries = {}
    top_k_percentage = self._loss_kwargs['top_k_percentage']
    deterministic = self._loss_kwargs['deterministic_top_k']
    rec_loss = self.rec_loss(seg, img, mask, top_k_percentage, deterministic)

    kl_dict = self.kl(seg, img)
    kl_sum = tf.reduce_sum(
        tf.stack([kl for _, kl in kl_dict.iteritems()], axis=-1))

    summaries['rec_loss_mean'] = rec_loss['mean']
    summaries['rec_loss_sum'] = rec_loss['sum']
    summaries['kl_sum'] = kl_sum
    for level, kl in kl_dict.iteritems():
      summaries['kl_{}'.format(level)] = kl

    # Set up a regular ELBO objective.
    if self._loss_kwargs['type'] == 'elbo':
      loss = rec_loss['sum'] + self._loss_kwargs['beta'] * kl_sum
      summaries['elbo_loss'] = loss

    # Set up a GECO objective (ELBO with a reconstruction constraint).
    elif self._loss_kwargs['type'] == 'geco':
      ma_rec_loss = self._moving_average(rec_loss['sum'])
      mask_sum_per_instance = tf.reduce_sum(rec_loss['mask'], axis=-1)
      num_valid_pixels = tf.reduce_mean(mask_sum_per_instance)
      reconstruction_threshold = self._loss_kwargs['kappa'] * num_valid_pixels

      rec_constraint = ma_rec_loss - reconstruction_threshold
      lagmul = self._lagmul(rec_constraint)
      loss = lagmul * rec_constraint + kl_sum

      summaries['geco_loss'] = loss
      summaries['ma_rec_loss_mean'] = ma_rec_loss / num_valid_pixels
      summaries['num_valid_pixels'] = num_valid_pixels
      summaries['lagmul'] = lagmul
    else:
      raise NotImplementedError('Loss type {} not implemeted!'.format(
          self._loss_kwargs['type']))

    return dict(supervised_loss=loss, summaries=summaries)

if __name__ == '__main__':
  hpu_net = HierarchicalProbUNet()