conv_transpose.py

# Copyright 2019 The Sonnet Authors. All Rights Reserved.
#
# 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.
# ============================================================================
"""Transpose convolutional module."""

from typing import Optional, Sequence, Union

import numpy as np
from sonnet.src import base
from sonnet.src import initializers
from sonnet.src import once
from sonnet.src import types
from sonnet.src import utils
import tensorflow as tf


def smart_concat(v1, v2):
  if isinstance(v1, tf.Tensor) or isinstance(v2, tf.Tensor):
    return tf.concat([v1, v2], 0)
  else:
    return v1 + v2


def smart_lambda(func, v1, v2):
  if isinstance(v1, tf.Tensor) or isinstance(v2, tf.Tensor):
    return func(v1, v2)
  else:
    return [func(x, y) for (x, y) in zip(v1, v2)]


class ConvNDTranspose(base.Module):
  """An N-dimensional transpose convolutional module.

  Attributes:
     w: Weight variable. Note is `None` until module is connected.
     b: Biases variable. Note is `None` until module is connected.
     input_shape: The input shape of the first set of inputs. Note is `None`
       until module is connected.
  """

  def __init__(self,
               num_spatial_dims: int,
               output_channels: int,
               kernel_shape: Union[int, Sequence[int]],
               output_shape: Optional[types.ShapeLike] = None,
               stride: Union[int, Sequence[int]] = 1,
               rate: Union[int, Sequence[int]] = 1,
               padding: str = "SAME",
               with_bias: bool = True,
               w_init: Optional[initializers.Initializer] = None,
               b_init: Optional[initializers.Initializer] = None,
               data_format: Optional[str] = None,
               name: Optional[str] = None):
    """Constructs a `ConvNDTranspose` module.

    Args:
      num_spatial_dims: Number of spatial dimensions of the input.
      output_channels: Number of output channels.
      kernel_shape: Sequence of integers (of length num_spatial_dims), or an
        integer representing kernel shape. `kernel_shape` will be expanded to
        define a kernel size in all dimensions.
      output_shape: Output shape of the spatial dimensions of a transpose
        convolution. Can be either an iterable of integers or a
        `TensorShape` of length `num_spatial_dims`. If a `None` value is given,
        a default shape is automatically calculated.
      stride: Sequence of integers (of length num_spatial_dims), or an integer.
        `stride` will be expanded to define stride in all dimensions.
      rate: Sequence of integers (of length num_spatial_dims), or integer that
        is used to define dilation rate in all dimensions. 1 corresponds to
        standard ND convolution, `rate > 1` corresponds to dilated convolution.
      padding: Padding algorithm, either "SAME" or "VALID".
      with_bias: Boolean, whether to include bias parameters. Default `True`.
      w_init: Optional initializer for the weights. By default the weights are
        initialized truncated random normal values with a standard deviation of
        `1 / sqrt(input_feature_size)`, which is commonly used when the
        inputs are zero centered (see https://arxiv.org/abs/1502.03167v3).
      b_init: Optional initializer for the bias. By default the bias is
        initialized to zero.
      data_format: The data format of the input.
      name: Name of the module.
    """
    super().__init__(name=name)

    if not 1 <= num_spatial_dims <= 3:
      raise ValueError(
          "We only support transpose convolution operations for "
          "num_spatial_dims=1, 2 or 3, received num_spatial_dims={}.".format(
              num_spatial_dims))
    self._num_spatial_dims = num_spatial_dims
    self._output_channels = output_channels
    self._kernel_shape = kernel_shape
    self._output_shape = output_shape
    self._stride = stride
    self._rate = rate

    if padding == "SAME" or padding == "VALID":
      self._padding = padding
    else:
      raise TypeError("ConvNDTranspose only takes string padding, please "
                      "provide either `SAME` or `VALID`.")
    self._data_format = data_format
    self._channel_index = utils.get_channel_index(data_format)
    self._with_bias = with_bias

    self._w_init = w_init
    if with_bias:
      self._b_init = b_init if b_init is not None else initializers.Zeros()
    elif b_init is not None:
      raise ValueError("When not using a bias the b_init must be None.")

  def __call__(self, inputs):
    self._initialize(inputs)

    if self._output_shape is None:
      output_shape = self._get_output_shape(inputs)
      if self._channel_index == 1:
        output_shape = smart_concat([self._output_channels], output_shape)
      else:
        output_shape = smart_concat(output_shape, [self._output_channels])
    else:
      output_shape = self._output_shape
    output_shape = smart_concat([tf.shape(inputs)[0]], output_shape)

    outputs = tf.nn.conv_transpose(
        input=inputs,
        filters=self.w,
        output_shape=output_shape,
        strides=self._stride,
        padding=self._padding,
        data_format=self._data_format,
        dilations=self._rate,
        name=None)
    if self._with_bias:
      outputs = tf.nn.bias_add(outputs, self.b, data_format=self._data_format)
    return outputs

  @once.once
  def _initialize(self, inputs):
    utils.assert_rank(inputs, self._num_spatial_dims + 2)
    self.input_channels = inputs.shape[self._channel_index]
    if self.input_channels is None:
      raise ValueError("The number of input channels must be known")
    self._dtype = inputs.dtype

    if self._output_shape is not None:
      if len(self._output_shape) != self._num_spatial_dims:
        raise ValueError(
            "The output_shape must be of length {} but instead was {}.".format(
                self._num_spatial_dims, len(self._output_shape)))
      if self._channel_index == 1:
        self._output_shape = [self._output_channels] + list(self._output_shape)
      else:
        self._output_shape = list(self._output_shape) + [self._output_channels]

    self.w = self._make_w()
    if self._with_bias:
      self.b = tf.Variable(
          self._b_init((self._output_channels,), self._dtype), name="b")

  def _make_w(self):
    """Makes and returns the variable representing the weight."""
    kernel_shape = utils.replicate(self._kernel_shape, self._num_spatial_dims,
                                   "kernel_shape")
    weight_shape = kernel_shape + (self._output_channels, self.input_channels)

    if self._w_init is None:
      # See https://arxiv.org/abs/1502.03167v3.
      fan_in_shape = kernel_shape + (self.input_channels,)
      stddev = 1 / np.sqrt(np.prod(fan_in_shape))
      self._w_init = initializers.TruncatedNormal(stddev=stddev)

    return tf.Variable(self._w_init(weight_shape, self._dtype), name="w")

  def _get_output_shape(self, inputs):
    input_shape = inputs.shape if inputs.shape.is_fully_defined() else tf.shape(
        inputs)

    if self._channel_index == 1:
      input_size = input_shape[2:]
    else:
      input_size = input_shape[1:-1]
    stride = utils.replicate(self._stride, self._num_spatial_dims, "stride")

    output_shape = smart_lambda(lambda x, y: x * y, input_size, stride)

    if self._padding == "VALID":
      kernel_shape = utils.replicate(self._kernel_shape, self._num_spatial_dims,
                                     "kernel_shape")
      rate = utils.replicate(self._rate, self._num_spatial_dims, "rate")
      effective_kernel_shape = [
          (shape - 1) * rate + 1 for (shape, rate) in zip(kernel_shape, rate)
      ]
      output_shape = smart_lambda(lambda x, y: x + y - 1, output_shape,
                                  effective_kernel_shape)

    return output_shape


class Conv1DTranspose(ConvNDTranspose):
  """A 1D transpose convolutional module."""

  def __init__(self,
               output_channels: int,
               kernel_shape: Union[int, Sequence[int]],
               output_shape: Optional[types.ShapeLike] = None,
               stride: Union[int, Sequence[int]] = 1,
               rate: Union[int, Sequence[int]] = 1,
               padding: str = "SAME",
               with_bias: bool = True,
               w_init: Optional[initializers.Initializer] = None,
               b_init: Optional[initializers.Initializer] = None,
               data_format: str = "NWC",
               name: Optional[str] = None):
    """Constructs a `Conv1DTranspose` module.

    Args:
      output_channels: Number of output channels.
      kernel_shape: Sequence of integers (of length 1), or an integer
        representing kernel shape. `kernel_shape` will be expanded to define a
        kernel size in all dimensions.
      output_shape: Output shape of the spatial dimensions of a transpose
        convolution. Can be either an integer or an iterable of integers or
        `Dimension`s, or a `TensorShape` (of length 1). If a `None` value is
        given, a default shape is automatically calculated.
      stride: Sequence of integers (of length 1), or an integer. `stride` will
        be expanded to define stride in all dimensions.
      rate: Sequence of integers (of length 1), or integer that is used to
        define dilation rate in all dimensions. 1 corresponds to standard 1D
        convolution, `rate > 1` corresponds to dilated convolution.
      padding: Padding algorithm, either "SAME" or "VALID".
      with_bias: Boolean, whether to include bias parameters. Default `True`.
      w_init: Optional initializer for the weights. By default the weights are
        initialized truncated random normal values with a standard deviation of
        `1 / sqrt(input_feature_size)`, which is commonly used when the
        inputs are zero centered (see https://arxiv.org/abs/1502.03167v3).
      b_init: Optional initializer for the bias. By default the bias is
        initialized to zero.
      data_format: The data format of the input.
      name: Name of the module.
    """
    super().__init__(
        num_spatial_dims=1,
        output_channels=output_channels,
        kernel_shape=kernel_shape,
        output_shape=output_shape,
        stride=stride,
        rate=rate,
        padding=padding,
        with_bias=with_bias,
        w_init=w_init,
        b_init=b_init,
        data_format=data_format,
        name=name)


class Conv2DTranspose(ConvNDTranspose):
  """A 2D transpose convolutional module."""

  def __init__(self,
               output_channels: int,
               kernel_shape: Union[int, Sequence[int]],
               output_shape: Optional[types.ShapeLike] = None,
               stride: Union[int, Sequence[int]] = 1,
               rate: Union[int, Sequence[int]] = 1,
               padding: str = "SAME",
               with_bias: bool = True,
               w_init: Optional[initializers.Initializer] = None,
               b_init: Optional[initializers.Initializer] = None,
               data_format: str = "NHWC",
               name: Optional[str] = None):
    """Constructs a `Conv2DTranspose` module.

    Args:
      output_channels: An integer, The number of output channels.
      kernel_shape: Sequence of integers (of length 2), or an integer
        representing kernel shape. `kernel_shape` will be expanded to define a
        kernel size in all dimensions.
      output_shape: Output shape of the spatial dimensions of a transpose
        convolution. Can be either an integer or an iterable of integers or
        `Dimension`s, or a `TensorShape` (of length 2). If a `None` value is
        given, a default shape is automatically calculated.
      stride: Sequence of integers (of length 2), or an integer. `stride` will
        be expanded to define stride in all dimensions.
      rate: Sequence of integers (of length 2), or integer that is used to
        define dilation rate in all dimensions. 1 corresponds to standard 2D
        convolution, `rate > 1` corresponds to dilated convolution.
      padding: Padding algorithm, either "SAME" or "VALID".
      with_bias: Boolean, whether to include bias parameters. Default `True`.
      w_init: Optional initializer for the weights. By default the weights are
        initialized truncated random normal values with a standard deviation of
        `1 / sqrt(input_feature_size)`, which is commonly used when the
        inputs are zero centered (see https://arxiv.org/abs/1502.03167v3).
      b_init: Optional initializer for the bias. By default the bias is
        initialized to zero.
      data_format: The data format of the input.
      name: Name of the module.
    """
    super().__init__(
        num_spatial_dims=2,
        output_channels=output_channels,
        kernel_shape=kernel_shape,
        output_shape=output_shape,
        stride=stride,
        rate=rate,
        padding=padding,
        with_bias=with_bias,
        w_init=w_init,
        b_init=b_init,
        data_format=data_format,
        name=name)


class Conv3DTranspose(ConvNDTranspose):
  """A 3D transpose convolutional module."""

  def __init__(self,
               output_channels: int,
               kernel_shape: Union[int, Sequence[int]],
               output_shape: Optional[types.ShapeLike] = None,
               stride: Union[int, Sequence[int]] = 1,
               rate: Union[int, Sequence[int]] = 1,
               padding: str = "SAME",
               with_bias: bool = True,
               w_init: Optional[initializers.Initializer] = None,
               b_init: Optional[initializers.Initializer] = None,
               data_format: str = "NDHWC",
               name: Optional[str] = None):
    """Constructs a `Conv3DTranspose` module.

    Args:
      output_channels: An integer, The number of output channels.
      kernel_shape: Sequence of integers (of length 3), or an integer
        representing kernel shape. `kernel_shape` will be expanded to define a
        kernel size in all dimensions.
      output_shape: Output shape of the spatial dimensions of a transpose
        convolution. Can be either an integer or an iterable of integers or
        `Dimension`s, or a `TensorShape` (of length 3). If a None value is
        given, a default shape is automatically calculated.
      stride: Sequence of integers (of length 3), or an integer. `stride` will
        be expanded to define stride in all dimensions.
      rate: Sequence of integers (of length 3), or integer that is used to
        define dilation rate in all dimensions. 1 corresponds to standard 3D
        convolution, `rate > 1` corresponds to dilated convolution.
      padding: Padding algorithm, either "SAME" or "VALID".
      with_bias: Boolean, whether to include bias parameters. Default `True`.
      w_init: Optional initializer for the weights. By default the weights are
        initialized truncated random normal values with a standard deviation of
        `1 / sqrt(input_feature_size)`, which is commonly used when the
        inputs are zero centered (see https://arxiv.org/abs/1502.03167v3).
      b_init: Optional initializer for the bias. By default the bias is
        initialized to zero.
      data_format: The data format of the input.
      name: Name of the module.
    """
    super().__init__(
        num_spatial_dims=3,
        output_channels=output_channels,
        kernel_shape=kernel_shape,
        output_shape=output_shape,
        stride=stride,
        rate=rate,
        padding=padding,
        with_bias=with_bias,
        w_init=w_init,
        b_init=b_init,
        data_format=data_format,
        name=name)