resnet.py

# Copyright 2023 The medical_research_foundations 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.

"""Contains definitions for the post-activation form of Residual Networks.

Residual networks (ResNets) were proposed in:
[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
    Deep Residual Learning for Image Recognition. arXiv:1512.03385
"""
from typing import Callable, Dict, Optional, Tuple

from absl import flags
from . import bit
import tensorflow.compat.v1 as tf
import typing_extensions

# pylint:disable=g-direct-tensorflow-import
from tensorflow.python.tpu import tpu_function
# pylint:enable=g-direct-tensorflow-import

FLAGS = flags.FLAGS
BATCH_NORM_EPSILON = 1e-5


class BlockFn(typing_extensions.Protocol):
  """Typing for block functions."""

  def __call__(
      self,
      inputs: tf.Tensor,
      filters: int,
      is_training: bool,
      strides: int,
      use_projection: bool = False,
      data_format: str = 'channels_first',
      dropblock_keep_prob: Optional[Tuple[float, float, float, float]] = None,
      dropblock_size: Optional[int] = None,
      global_bn: bool = True,
      batch_norm_decay: float = 0.9,
  ) -> tf.Tensor:
    pass


class BatchNormalization(tf.layers.BatchNormalization):
  """Batch Normalization layer that supports cross replica computation on TPU.

  This class extends the keras.BatchNormalization implementation by supporting
  cross replica means and variances. The base class implementation only computes
  moments based on mini-batch per replica (TPU core).

  For detailed information of arguments and implementation, refer to:
  https://www.tensorflow.org/api_docs/python/tf/keras/layers/BatchNormalization
  """

  def __init__(self, fused=False, **kwargs):
    """Builds the batch normalization layer.

    Arguments:
      fused: If `False`, use the system recommended implementation. Only support
        `False` in the current implementation.
      **kwargs: input augments that are forwarded to
        tf.layers.BatchNormalization.
    """
    if fused in (True, None):
      raise ValueError('The TPU version of BatchNormalization does not support '
                       'fused=True.')
    super(BatchNormalization, self).__init__(fused=fused, **kwargs)

  def _cross_replica_average(self, t):
    """Calculates the average value of input tensor across TPU replicas."""
    num_shards = tpu_function.get_tpu_context().number_of_shards
    return tf.tpu.cross_replica_sum(t) / tf.cast(num_shards, t.dtype)

  def _moments(self, inputs, reduction_axes, keep_dims, mask=None):
    """Compute the mean and variance: it overrides the original _moments."""
    shard_mean, shard_variance = super(BatchNormalization, self)._moments(
        inputs, reduction_axes, keep_dims=keep_dims, mask=mask)

    num_shards = tpu_function.get_tpu_context().number_of_shards
    if num_shards and num_shards > 1:
      # Each group has multiple replicas: here we compute group mean/variance by
      # aggregating per-replica mean/variance.
      group_mean = self._cross_replica_average(shard_mean)
      group_variance = self._cross_replica_average(shard_variance)

      # Group variance needs to also include the difference between shard_mean
      # and group_mean.
      mean_distance = tf.square(group_mean - shard_mean)
      group_variance += self._cross_replica_average(mean_distance)
      return (group_mean, group_variance)
    else:
      return (shard_mean, shard_variance)


def batch_norm_relu(
    inputs,
    is_training,
    relu=True,
    init_zero=False,
    center=True,
    scale=True,
    data_format='channels_first',
    global_bn=True,
    batch_norm_decay=0.9,
):
  """Performs a batch normalization followed by a ReLU.

  Args:
    inputs: `Tensor` of shape `[batch, channels, ...]`.
    is_training: `bool` for whether the model is training.
    relu: `bool` if False, omits the ReLU operation.
    init_zero: `bool` if True, initializes scale parameter of batch
      normalization with 0 instead of 1 (default).
    center: `bool` whether to add learnable bias factor.
    scale: `bool` whether to add learnable scaling factor.
    data_format: `str` either "channels_first" for `[batch, channels, height,
      width]` or "channels_last for `[batch, height, width, channels]`.
    global_bn: `bool` whether to aggregate BN statistics across distributed
      cores.
    batch_norm_decay: `float` batch norm decay parameter.

  Returns:
    A normalized `Tensor` with the same `data_format`.
  """
  if init_zero:
    gamma_initializer = tf.zeros_initializer()
  else:
    gamma_initializer = tf.ones_initializer()

  if data_format == 'channels_first':
    axis = 1
  else:
    axis = 3

  if global_bn:
    bn_foo = BatchNormalization(
        axis=axis,
        momentum=batch_norm_decay,
        epsilon=BATCH_NORM_EPSILON,
        center=center,
        scale=scale,
        fused=False,
        gamma_initializer=gamma_initializer,
    )
    inputs = bn_foo(inputs, training=is_training)
  else:
    inputs = tf.layers.batch_normalization(
        inputs=inputs,
        axis=axis,
        momentum=batch_norm_decay,
        epsilon=BATCH_NORM_EPSILON,
        center=center,
        scale=scale,
        training=is_training,
        fused=True,
        gamma_initializer=gamma_initializer,
    )

  if relu:
    inputs = tf.nn.relu(inputs)
  return inputs


def dropblock(
    net, is_training, keep_prob, dropblock_size, data_format='channels_first'
):
  """DropBlock: a regularization method for convolutional neural networks.

  DropBlock is a form of structured dropout, where units in a contiguous
  region of a feature map are dropped together. DropBlock works better than
  dropout on convolutional layers due to the fact that activation units in
  convolutional layers are spatially correlated.
  See https://arxiv.org/pdf/1810.12890.pdf for details.

  Args:
    net: `Tensor` input tensor.
    is_training: `bool` for whether the model is training.
    keep_prob: `float` or `Tensor` keep_prob parameter of DropBlock. "None"
      means no DropBlock.
    dropblock_size: `int` size of blocks to be dropped by DropBlock.
    data_format: `str` either "channels_first" for `[batch, channels, height,
      width]` or "channels_last for `[batch, height, width, channels]`.

  Returns:
      A version of input tensor with DropBlock applied.
  Raises:
      if width and height of the input tensor are not equal.
  """

  if not is_training or keep_prob is None:
    return net

  tf.logging.info('Applying DropBlock: dropblock_size {}, net.shape {}'.format(
      dropblock_size, net.shape))

  if data_format == 'channels_last':
    _, width, height, _ = net.get_shape().as_list()
  else:
    _, _, width, height = net.get_shape().as_list()
  if width != height:
    raise ValueError('Input tensor with width!=height is not supported.')

  dropblock_size = min(dropblock_size, width)
  # seed_drop_rate is the gamma parameter of DropBlcok.
  seed_drop_rate = (1.0 - keep_prob) * width**2 / dropblock_size**2 / (
      width - dropblock_size + 1)**2

  # Forces the block to be inside the feature map.
  w_i, h_i = tf.meshgrid(tf.range(width), tf.range(width))
  valid_block_center = tf.logical_and(
      tf.logical_and(w_i >= int(dropblock_size // 2),
                     w_i < width - (dropblock_size - 1) // 2),
      tf.logical_and(h_i >= int(dropblock_size // 2),
                     h_i < width - (dropblock_size - 1) // 2))

  valid_block_center = tf.expand_dims(valid_block_center, 0)
  valid_block_center = tf.expand_dims(
      valid_block_center, -1 if data_format == 'channels_last' else 0)

  randnoise = tf.random_uniform(net.shape, dtype=tf.float32)
  block_pattern = (1 - tf.cast(valid_block_center, dtype=tf.float32) + tf.cast(
      (1 - seed_drop_rate), dtype=tf.float32) + randnoise) >= 1
  block_pattern = tf.cast(block_pattern, dtype=tf.float32)

  if dropblock_size == width:
    block_pattern = tf.reduce_min(
        block_pattern,
        axis=[1, 2] if data_format == 'channels_last' else [2, 3],
        keepdims=True)
  else:
    if data_format == 'channels_last':
      ksize = [1, dropblock_size, dropblock_size, 1]
    else:
      ksize = [1, 1, dropblock_size, dropblock_size]
    block_pattern = -tf.nn.max_pool(
        -block_pattern, ksize=ksize, strides=[1, 1, 1, 1], padding='SAME',
        data_format='NHWC' if data_format == 'channels_last' else 'NCHW')

  percent_ones = tf.cast(tf.reduce_sum((block_pattern)), tf.float32) / tf.cast(
      tf.size(block_pattern), tf.float32)

  net = net / tf.cast(percent_ones, net.dtype) * tf.cast(
      block_pattern, net.dtype)
  return net


def fixed_padding(inputs, kernel_size, data_format='channels_first'):
  """Pads the input along the spatial dimensions independently of input size.

  Args:
    inputs: `Tensor` of size `[batch, channels, height, width]` or `[batch,
      height, width, channels]` depending on `data_format`.
    kernel_size: `int` kernel size to be used for `conv2d` or max_pool2d`
      operations. Should be a positive integer.
    data_format: `str` either "channels_first" for `[batch, channels, height,
      width]` or "channels_last for `[batch, height, width, channels]`.

  Returns:
    A padded `Tensor` of the same `data_format` with size either intact
    (if `kernel_size == 1`) or padded (if `kernel_size > 1`).
  """
  pad_total = kernel_size - 1
  pad_beg = pad_total // 2
  pad_end = pad_total - pad_beg
  if data_format == 'channels_first':
    padded_inputs = tf.pad(inputs, [[0, 0], [0, 0],
                                    [pad_beg, pad_end], [pad_beg, pad_end]])
  else:
    padded_inputs = tf.pad(inputs, [[0, 0], [pad_beg, pad_end],
                                    [pad_beg, pad_end], [0, 0]])

  return padded_inputs


def conv2d_fixed_padding(
    inputs, filters, kernel_size, strides, data_format='channels_first'
):
  """Strided 2-D convolution with explicit padding.

  The padding is consistent and is based only on `kernel_size`, not on the
  dimensions of `inputs` (as opposed to using `tf.layers.conv2d` alone).

  Args:
    inputs: `Tensor` of size `[batch, channels, height_in, width_in]`.
    filters: `int` number of filters in the convolution.
    kernel_size: `int` size of the kernel to be used in the convolution.
    strides: `int` strides of the convolution.
    data_format: `str` either "channels_first" for `[batch, channels, height,
      width]` or "channels_last for `[batch, height, width, channels]`.

  Returns:
    A `Tensor` of shape `[batch, filters, height_out, width_out]`.
  """
  if strides > 1:
    inputs = fixed_padding(inputs, kernel_size, data_format=data_format)

  return tf.layers.conv2d(
      inputs=inputs, filters=filters, kernel_size=kernel_size, strides=strides,
      padding=('SAME' if strides == 1 else 'VALID'), use_bias=False,
      kernel_initializer=tf.variance_scaling_initializer(),
      data_format=data_format)


def residual_block(
    inputs,
    filters,
    is_training,
    strides,
    use_projection=False,
    data_format='channels_first',
    dropblock_keep_prob=None,
    dropblock_size=None,
    global_bn=True,
    batch_norm_decay=0.9,
):
  """Standard building block for residual networks with BN after convolutions.

  Args:
    inputs: `Tensor` of size `[batch, channels, height, width]`.
    filters: `int` number of filters for the first two convolutions. Note that
      the third and final convolution will use 4 times as many filters.
    is_training: `bool` for whether the model is in training.
    strides: `int` block stride. If greater than 1, this block will ultimately
      downsample the input.
    use_projection: `bool` for whether this block should use a projection
      shortcut (versus the default identity shortcut). This is usually `True`
      for the first block of a block group, which may change the number of
      filters and the resolution.
    data_format: `str` either "channels_first" for `[batch, channels, height,
      width]` or "channels_last for `[batch, height, width, channels]`.
    dropblock_keep_prob: unused; needed to give method same signature as other
      blocks
    dropblock_size: unused; needed to give method same signature as other
      blocks.
    global_bn: `bool` whether to aggregate BN statistics across distributed
      cores.
    batch_norm_decay: `float` batch norm decay parameter.

  Returns:
    The output `Tensor` of the block.
  """
  del dropblock_keep_prob
  del dropblock_size
  shortcut = inputs
  if use_projection:
    # Projection shortcut in first layer to match filters and strides
    shortcut = conv2d_fixed_padding(
        inputs=inputs,
        filters=filters,
        kernel_size=1,
        strides=strides,
        data_format=data_format,
    )
    shortcut = batch_norm_relu(
        shortcut,
        is_training,
        relu=False,
        data_format=data_format,
        global_bn=global_bn,
        batch_norm_decay=batch_norm_decay,
    )

  inputs = conv2d_fixed_padding(
      inputs=inputs, filters=filters, kernel_size=3, strides=strides,
      data_format=data_format)
  inputs = batch_norm_relu(
      inputs,
      is_training,
      data_format=data_format,
      global_bn=global_bn,
      batch_norm_decay=batch_norm_decay,
  )

  inputs = conv2d_fixed_padding(
      inputs=inputs, filters=filters, kernel_size=3, strides=1,
      data_format=data_format)
  inputs = batch_norm_relu(
      inputs,
      is_training,
      relu=False,
      init_zero=True,
      data_format=data_format,
      global_bn=global_bn,
      batch_norm_decay=batch_norm_decay,
  )

  return tf.nn.relu(inputs + shortcut)


def bottleneck_block(
    inputs,
    filters,
    is_training,
    strides,
    use_projection=False,
    data_format='channels_first',
    dropblock_keep_prob=None,
    dropblock_size=None,
    global_bn=True,
    batch_norm_decay=0.9,
):
  """Bottleneck block variant for residual networks with BN after convolutions.

  Args:
    inputs: `Tensor` of size `[batch, channels, height, width]`.
    filters: `int` number of filters for the first two convolutions. Note that
      the third and final convolution will use 4 times as many filters.
    is_training: `bool` for whether the model is in training.
    strides: `int` block stride. If greater than 1, this block will ultimately
      downsample the input.
    use_projection: `bool` for whether this block should use a projection
      shortcut (versus the default identity shortcut). This is usually `True`
      for the first block of a block group, which may change the number of
      filters and the resolution.
    data_format: `str` either "channels_first" for `[batch, channels, height,
      width]` or "channels_last for `[batch, height, width, channels]`.
    dropblock_keep_prob: `float` or `Tensor` keep_prob parameter of DropBlock.
      "None" means no DropBlock.
    dropblock_size: `int` size parameter of DropBlock. Will not be used if
      dropblock_keep_prob is "None".
    global_bn: `bool` whether to aggregate BN statistics across distributed
      cores.
    batch_norm_decay: `float` batch norm decay parameter.

  Returns:
    The output `Tensor` of the block.
  """
  shortcut = inputs
  if use_projection:
    # Projection shortcut only in first block within a group. Bottleneck blocks
    # end with 4 times the number of filters.
    filters_out = 4 * filters
    shortcut = conv2d_fixed_padding(
        inputs=inputs,
        filters=filters_out,
        kernel_size=1,
        strides=strides,
        data_format=data_format,
    )
    shortcut = batch_norm_relu(
        shortcut,
        is_training,
        relu=False,
        data_format=data_format,
        global_bn=global_bn,
        batch_norm_decay=batch_norm_decay,
    )
  shortcut = dropblock(
      shortcut, is_training=is_training, data_format=data_format,
      keep_prob=dropblock_keep_prob, dropblock_size=dropblock_size)

  inputs = conv2d_fixed_padding(
      inputs=inputs, filters=filters, kernel_size=1, strides=1,
      data_format=data_format)
  inputs = batch_norm_relu(
      inputs,
      is_training,
      data_format=data_format,
      global_bn=global_bn,
      batch_norm_decay=batch_norm_decay,
  )
  inputs = dropblock(
      inputs, is_training=is_training, data_format=data_format,
      keep_prob=dropblock_keep_prob, dropblock_size=dropblock_size)

  inputs = conv2d_fixed_padding(
      inputs=inputs,
      filters=filters,
      kernel_size=3,
      strides=strides,
      data_format=data_format,
  )
  inputs = batch_norm_relu(
      inputs,
      is_training,
      data_format=data_format,
      global_bn=global_bn,
      batch_norm_decay=batch_norm_decay,
  )
  inputs = dropblock(
      inputs, is_training=is_training, data_format=data_format,
      keep_prob=dropblock_keep_prob, dropblock_size=dropblock_size)

  inputs = conv2d_fixed_padding(
      inputs=inputs, filters=4 * filters, kernel_size=1, strides=1,
      data_format=data_format)
  inputs = batch_norm_relu(
      inputs,
      is_training,
      relu=False,
      init_zero=True,
      data_format=data_format,
      global_bn=global_bn,
      batch_norm_decay=batch_norm_decay,
  )
  inputs = dropblock(
      inputs, is_training=is_training, data_format=data_format,
      keep_prob=dropblock_keep_prob, dropblock_size=dropblock_size)

  return tf.nn.relu(inputs + shortcut)


def block_group(
    inputs,
    filters,
    block_fn,
    blocks,
    strides,
    is_training,
    name,
    data_format='channels_first',
    dropblock_keep_prob=None,
    dropblock_size=None,
):
  """Creates one group of blocks for the ResNet model.

  Args:
    inputs: `Tensor` of size `[batch, channels, height, width]`.
    filters: `int` number of filters for the first convolution of the layer.
    block_fn: `function` for the block to use within the model
    blocks: `int` number of blocks contained in the layer.
    strides: `int` stride to use for the first convolution of the layer. If
      greater than 1, this layer will downsample the input.
    is_training: `bool` for whether the model is training.
    name: `str`name for the Tensor output of the block layer.
    data_format: `str` either "channels_first" for `[batch, channels, height,
      width]` or "channels_last for `[batch, height, width, channels]`.
    dropblock_keep_prob: `float` or `Tensor` keep_prob parameter of DropBlock.
      "None" means no DropBlock.
    dropblock_size: `int` size parameter of DropBlock. Will not be used if
      dropblock_keep_prob is "None".

  Returns:
    The output `Tensor` of the block layer.
  """
  # Only the first block per block_group uses projection shortcut and strides.
  inputs = block_fn(inputs, filters, is_training, strides,
                    use_projection=True, data_format=data_format,
                    dropblock_keep_prob=dropblock_keep_prob,
                    dropblock_size=dropblock_size)

  for _ in range(1, blocks):
    inputs = block_fn(inputs, filters, is_training, 1,
                      data_format=data_format,
                      dropblock_keep_prob=dropblock_keep_prob,
                      dropblock_size=dropblock_size)

  return tf.identity(inputs, name)


def _resnet_v1_generator(
    block_fn: BlockFn,
    layers: Tuple[int, int, int, int],
    width_multiplier: int,
    cifar_stem: bool = False,
    data_format: str = 'channels_last',
    dropblock_keep_probs: Optional[Tuple[float, float, float, float]] = None,
    dropblock_size: Optional[int] = None,
    train_mode: str = 'pretrain',
    fine_tune_after_block: int = -1,
    global_bn: bool = True,
    batch_norm_decay: float = 0.9,
) -> Callable[[tf.Tensor, bool], Dict[str, tf.Tensor]]:
  """Generator for ResNet v1 models.

  Args:
    block_fn: `function` for the block to use within the model. Either
      `residual_block` or `bottleneck_block`.
    layers: list of 4 `int`s denoting the number of blocks to include in each of
      the 4 block groups. Each group consists of blocks that take inputs of the
      same resolution.
    width_multiplier: `int` width multiplier for network.
    cifar_stem: `bool` If True, use a 3x3 conv without strides or pooling as
      stem.
    data_format: `str` either "channels_first" for `[batch, channels, height,
      width]` or "channels_last for `[batch, height, width, channels]`.
    dropblock_keep_probs: `list` of 4 elements denoting keep_prob of DropBlock
      for each block group. None indicates no DropBlock for the corresponding
      block group.
    dropblock_size: `int`: size parameter of DropBlock.
    train_mode: `str` either "pretrain" or "finetune".
    fine_tune_after_block: `int` the layers after which block that we will
      fine-tune. -1 means fine-tuning everything. 0 means fine-tuning after stem
      block. 4 means fine-tuning just the linear head.
    global_bn: `bool` whether to aggregate BN statistics across distributed
      cores.
    batch_norm_decay: `float` batch norm decay parameter.

  Returns:
    Model `function` that takes in `inputs` and `is_training` and returns the
    output `Tensor` of the ResNet model.

  Raises:
    if dropblock_keep_probs is not 'None' or a list with len 4.
  """
  if dropblock_keep_probs is None:
    dropblock_keep_probs = [None] * 4
  if not isinstance(dropblock_keep_probs,
                    list) or len(dropblock_keep_probs) != 4:
    raise ValueError('dropblock_keep_probs is not valid:', dropblock_keep_probs)

  def model(inputs: tf.Tensor, is_training: bool) -> Dict[str, tf.Tensor]:
    """Creation of the model graph."""
    outputs = {}
    if cifar_stem:
      inputs = conv2d_fixed_padding(
          inputs=inputs, filters=64 * width_multiplier, kernel_size=3,
          strides=1, data_format=data_format)
      inputs = tf.identity(inputs, 'initial_conv')
      inputs = batch_norm_relu(
          inputs,
          is_training,
          data_format=data_format,
          global_bn=global_bn,
          batch_norm_decay=batch_norm_decay,
      )
      inputs = tf.identity(inputs, 'initial_max_pool')
    else:
      inputs = conv2d_fixed_padding(
          inputs=inputs,
          filters=64 * width_multiplier,
          kernel_size=7,
          strides=2,
          data_format=data_format,
      )
      inputs = tf.identity(inputs, 'initial_conv')
      inputs = batch_norm_relu(
          inputs,
          is_training,
          data_format=data_format,
          global_bn=global_bn,
          batch_norm_decay=batch_norm_decay,
      )

      inputs = tf.layers.max_pooling2d(
          inputs=inputs, pool_size=3, strides=2, padding='SAME',
          data_format=data_format)
      inputs = tf.identity(inputs, 'initial_max_pool')

    def filter_trainable_variables(trainable_variables, after_block):
      """Add new trainable variables for the immediate precedent block."""
      if after_block == 0:
        trainable_variables[after_block] = tf.trainable_variables()
      else:
        trainable_variables[after_block] = []
        for var in tf.trainable_variables():
          to_keep = True
          for j in range(after_block):
            if var in trainable_variables[j]:
              to_keep = False
              break
          if to_keep:
            trainable_variables[after_block].append(var)

    def add_to_collection(trainable_variables, prefix):
      """Put variables into graph collection."""
      for after_block, variables in trainable_variables.items():
        collection = prefix + str(after_block)
        for var in variables:
          tf.add_to_collection(collection, var)

    trainable_variables = {}
    filter_trainable_variables(trainable_variables, after_block=0)
    if train_mode == 'finetune' and fine_tune_after_block == 0:
      inputs = tf.stop_gradient(inputs)

    inputs = block_group(
        inputs=inputs, filters=64 * width_multiplier, block_fn=block_fn,
        blocks=layers[0], strides=1, is_training=is_training,
        name='block_group1', data_format=data_format,
        dropblock_keep_prob=dropblock_keep_probs[0],
        dropblock_size=dropblock_size)
    outputs['block_group1'] = inputs

    filter_trainable_variables(trainable_variables, after_block=1)
    if train_mode == 'finetune' and fine_tune_after_block == 1:
      inputs = tf.stop_gradient(inputs)

    inputs = block_group(
        inputs=inputs, filters=128 * width_multiplier, block_fn=block_fn,
        blocks=layers[1], strides=2, is_training=is_training,
        name='block_group2', data_format=data_format,
        dropblock_keep_prob=dropblock_keep_probs[1],
        dropblock_size=dropblock_size)
    outputs['block_group2'] = inputs

    filter_trainable_variables(trainable_variables, after_block=2)
    if train_mode == 'finetune' and fine_tune_after_block == 2:
      inputs = tf.stop_gradient(inputs)

    inputs = block_group(
        inputs=inputs, filters=256 * width_multiplier, block_fn=block_fn,
        blocks=layers[2], strides=2, is_training=is_training,
        name='block_group3', data_format=data_format,
        dropblock_keep_prob=dropblock_keep_probs[2],
        dropblock_size=dropblock_size)
    outputs['block_group3'] = inputs

    filter_trainable_variables(trainable_variables, after_block=3)
    if train_mode == 'finetune' and fine_tune_after_block == 3:
      inputs = tf.stop_gradient(inputs)

    inputs = block_group(
        inputs=inputs, filters=512 * width_multiplier, block_fn=block_fn,
        blocks=layers[3], strides=2, is_training=is_training,
        name='block_group4', data_format=data_format,
        dropblock_keep_prob=dropblock_keep_probs[3],
        dropblock_size=dropblock_size)
    outputs['block_group4'] = inputs

    filter_trainable_variables(trainable_variables, after_block=4)
    if train_mode == 'finetune' and fine_tune_after_block == 4:
      inputs = tf.stop_gradient(inputs)

    # The activation is 7x7 so this is a global average pool.
    # TODO: reduce_mean will be faster.
    pool_size = (inputs.shape[1], inputs.shape[2])
    inputs = tf.layers.average_pooling2d(
        inputs=inputs,
        pool_size=pool_size,
        strides=1,
        padding='VALID',
        data_format=data_format,
    )
    inputs = tf.identity(inputs, 'final_avg_pool')
    inputs = tf.squeeze(inputs, (1, 2))
    outputs['final_avg_pool'] = inputs

    # filter_trainable_variables(trainable_variables, after_block=5)
    add_to_collection(trainable_variables, 'trainable_variables_inblock_')

    return outputs

  return model


def resnet_v2(
    depth: int, width_multiplier: int, verify_input_range: bool = False
) -> Callable[[tf.Tensor, bool], Dict[str, tf.Tensor]]:
  """Returns the ResNet-v2/SimCLR+BiT model for a given architecture."""
  bit_model_archs = [(50, 1), (50, 3), (101, 1), (101, 3), (152, 2), (152, 4)]
  assert (depth, width_multiplier) in bit_model_archs, (
      'There is no SimCLR+BiT model architecture for the requested model depth '
      'and width multiplier: ({},{}). Valid model combinations are: {}'.format(
          depth, width_multiplier, bit_model_archs
      )
  )
  model_name = f'ResNet-{depth}x{width_multiplier}'

  def model(inputs: tf.Tensor, is_training: bool):
    del is_training  # No difference between train/eval graph
    _, endpoints = bit.bit_embedding(
        inputs,
        model_name=model_name,
        trainable=True,
        verify_input_range=verify_input_range,
    )
    return endpoints

  return model


def resnet_v1(
    resnet_depth: int,
    width_multiplier: int,
    cifar_stem: bool = False,
    data_format: str = 'channels_last',
    dropblock_keep_probs: Optional[Tuple[float, float, float, float]] = None,
    dropblock_size: Optional[int] = None,
    train_mode: str = 'pretrain',
    fine_tune_after_block: int = -1,
    global_bn: bool = True,
    batch_norm_decay: float = 0.9,
) -> Callable[[tf.Tensor, bool], Dict[str, tf.Tensor]]:
  """Returns the ResNet-v1 model for a given size and number of output classes."""

  model_params = {
      18: {'block': residual_block, 'layers': [2, 2, 2, 2]},
      34: {'block': residual_block, 'layers': [3, 4, 6, 3]},
      50: {'block': bottleneck_block, 'layers': [3, 4, 6, 3]},
      101: {'block': bottleneck_block, 'layers': [3, 4, 23, 3]},
      152: {'block': bottleneck_block, 'layers': [3, 8, 36, 3]},
      200: {'block': bottleneck_block, 'layers': [3, 24, 36, 3]}
  }

  if resnet_depth not in model_params:
    raise ValueError('Not a valid resnet_depth:', resnet_depth)

  params = model_params[resnet_depth]
  return _resnet_v1_generator(
      params['block'],
      tuple(params['layers']),
      width_multiplier,
      cifar_stem=cifar_stem,
      dropblock_keep_probs=dropblock_keep_probs,
      dropblock_size=dropblock_size,
      data_format=data_format,
      train_mode=train_mode,
      fine_tune_after_block=fine_tune_after_block,
      global_bn=global_bn,
      batch_norm_decay=batch_norm_decay,
  )