imaginaire/layers/activation_norm.py

# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out LICENSE.md
# flake8: noqa E722
from types import SimpleNamespace

import torch

try:
    from torch.nn import SyncBatchNorm
except ImportError:
    from torch.nn import BatchNorm2d as SyncBatchNorm
from torch import nn
from torch.nn import functional as F
from .conv import LinearBlock, Conv2dBlock, HyperConv2d, PartialConv2dBlock
from .misc import PartialSequential, ApplyNoise


class AdaptiveNorm(nn.Module):
    r"""Adaptive normalization layer. The layer first normalizes the input, then
    performs an affine transformation using parameters computed from the
    conditional inputs.

    Args:
        num_features (int): Number of channels in the input tensor.
        cond_dims (int): Number of channels in the conditional inputs.
        weight_norm_type (str): Type of weight normalization.
            ``'none'``, ``'spectral'``, ``'weight'``, or ``'weight_demod'``.
        projection (bool): If ``True``, project the conditional input to gamma
            and beta using a fully connected layer, otherwise directly use
            the conditional input as gamma and beta.
        projection_bias (bool) If ``True``, use bias in the fully connected
            projection layer.
        separate_projection (bool): If ``True``, we will use two different
            layers for gamma and beta. Otherwise, we will use one layer. It
            matters only if you apply any weight norms to this layer.
        input_dim (int): Number of dimensions of the input tensor.
        activation_norm_type (str):
            Type of activation normalization.
            ``'none'``, ``'instance'``, ``'batch'``, ``'sync_batch'``,
            ``'layer'``,  ``'layer_2d'``, ``'group'``, ``'adaptive'``,
            ``'spatially_adaptive'`` or ``'hyper_spatially_adaptive'``.
        activation_norm_params (obj, optional, default=None):
            Parameters of activation normalization.
            If not ``None``, ``activation_norm_params.__dict__`` will be used as
            keyword arguments when initializing activation normalization.
    """

    def __init__(self, num_features, cond_dims, weight_norm_type='',
                 projection=True,
                 projection_bias=True,
                 separate_projection=False,
                 input_dim=2,
                 activation_norm_type='instance',
                 activation_norm_params=None,
                 apply_noise=False,
                 add_bias=True,
                 input_scale=1.0,
                 init_gain=1.0):
        super().__init__()
        if activation_norm_params is None:
            activation_norm_params = SimpleNamespace(affine=False)
        self.norm = get_activation_norm_layer(num_features,
                                              activation_norm_type,
                                              input_dim,
                                              **vars(activation_norm_params))
        if apply_noise:
            self.noise_layer = ApplyNoise()
        else:
            self.noise_layer = None

        if projection:
            if separate_projection:
                self.fc_gamma = \
                    LinearBlock(cond_dims, num_features,
                                weight_norm_type=weight_norm_type,
                                bias=projection_bias)
                self.fc_beta = \
                    LinearBlock(cond_dims, num_features,
                                weight_norm_type=weight_norm_type,
                                bias=projection_bias)
            else:
                self.fc = LinearBlock(cond_dims, num_features * 2,
                                      weight_norm_type=weight_norm_type,
                                      bias=projection_bias)

        self.projection = projection
        self.separate_projection = separate_projection
        self.input_scale = input_scale
        self.add_bias = add_bias
        self.conditional = True
        self.init_gain = init_gain

    def forward(self, x, y, noise=None, **_kwargs):
        r"""Adaptive Normalization forward.

        Args:
            x (N x C1 x * tensor): Input tensor.
            y (N x C2 tensor): Conditional information.
        Returns:
            out (N x C1 x * tensor): Output tensor.
        """
        y = y * self.input_scale
        if self.projection:
            if self.separate_projection:
                gamma = self.fc_gamma(y)
                beta = self.fc_beta(y)
                for _ in range(x.dim() - gamma.dim()):
                    gamma = gamma.unsqueeze(-1)
                    beta = beta.unsqueeze(-1)
            else:
                y = self.fc(y)
                for _ in range(x.dim() - y.dim()):
                    y = y.unsqueeze(-1)
                gamma, beta = y.chunk(2, 1)
        else:
            for _ in range(x.dim() - y.dim()):
                y = y.unsqueeze(-1)
            gamma, beta = y.chunk(2, 1)
        if self.norm is not None:
            x = self.norm(x)
        if self.noise_layer is not None:
            x = self.noise_layer(x, noise=noise)
        if self.add_bias:
            x = torch.addcmul(beta, x, 1 + gamma)
            return x
        else:
            return x * (1 + gamma), beta.squeeze(3).squeeze(2)


class SpatiallyAdaptiveNorm(nn.Module):
    r"""Spatially Adaptive Normalization (SPADE) initialization.

    Args:
        num_features (int) : Number of channels in the input tensor.
        cond_dims (int or list of int) : List of numbers of channels
            in the input.
        num_filters (int): Number of filters in SPADE.
        kernel_size (int): Kernel size of the convolutional filters in
            the SPADE layer.
        weight_norm_type (str): Type of weight normalization.
            ``'none'``, ``'spectral'``, or ``'weight'``.
        separate_projection (bool): If ``True``, we will use two different
            layers for gamma and beta. Otherwise, we will use one layer. It
            matters only if you apply any weight norms to this layer.
        activation_norm_type (str):
            Type of activation normalization.
            ``'none'``, ``'instance'``, ``'batch'``, ``'sync_batch'``,
            ``'layer'``,  ``'layer_2d'``, ``'group'``.
        activation_norm_params (obj, optional, default=None):
            Parameters of activation normalization.
            If not ``None``, ``activation_norm_params.__dict__`` will be used as
            keyword arguments when initializing activation normalization.
    """

    def __init__(self,
                 num_features,
                 cond_dims,
                 num_filters=128,
                 kernel_size=3,
                 weight_norm_type='',
                 separate_projection=False,
                 activation_norm_type='sync_batch',
                 activation_norm_params=None,
                 bias_only=False,
                 partial=False,
                 interpolation='nearest'):
        super().__init__()
        if activation_norm_params is None:
            activation_norm_params = SimpleNamespace(affine=False)
        padding = kernel_size // 2
        self.separate_projection = separate_projection
        self.mlps = nn.ModuleList()
        self.gammas = nn.ModuleList()
        self.betas = nn.ModuleList()
        self.bias_only = bias_only
        self.interpolation = interpolation

        # Make cond_dims a list.
        if type(cond_dims) != list:
            cond_dims = [cond_dims]

        # Make num_filters a list.
        if not isinstance(num_filters, list):
            num_filters = [num_filters] * len(cond_dims)
        else:
            assert len(num_filters) >= len(cond_dims)

        # Make partial a list.
        if not isinstance(partial, list):
            partial = [partial] * len(cond_dims)
        else:
            assert len(partial) >= len(cond_dims)

        for i, cond_dim in enumerate(cond_dims):
            mlp = []
            conv_block = PartialConv2dBlock if partial[i] else Conv2dBlock
            sequential = PartialSequential if partial[i] else nn.Sequential

            if num_filters[i] > 0:
                mlp += [conv_block(cond_dim,
                                   num_filters[i],
                                   kernel_size,
                                   padding=padding,
                                   weight_norm_type=weight_norm_type,
                                   nonlinearity='relu')]
            mlp_ch = cond_dim if num_filters[i] == 0 else num_filters[i]

            if self.separate_projection:
                if partial[i]:
                    raise NotImplementedError(
                        'Separate projection not yet implemented for ' +
                        'partial conv')
                self.mlps.append(nn.Sequential(*mlp))
                self.gammas.append(
                    conv_block(mlp_ch, num_features,
                               kernel_size,
                               padding=padding,
                               weight_norm_type=weight_norm_type))
                self.betas.append(
                    conv_block(mlp_ch, num_features,
                               kernel_size,
                               padding=padding,
                               weight_norm_type=weight_norm_type))
            else:
                mlp += [conv_block(mlp_ch, num_features * 2, kernel_size,
                                   padding=padding,
                                   weight_norm_type=weight_norm_type)]
                self.mlps.append(sequential(*mlp))

        self.norm = get_activation_norm_layer(num_features,
                                              activation_norm_type,
                                              2,
                                              **vars(activation_norm_params))
        self.conditional = True

    def forward(self, x, *cond_inputs, **_kwargs):
        r"""Spatially Adaptive Normalization (SPADE) forward.

        Args:
            x (N x C1 x H x W tensor) : Input tensor.
            cond_inputs (list of tensors) : Conditional maps for SPADE.
        Returns:
            output (4D tensor) : Output tensor.
        """
        output = self.norm(x) if self.norm is not None else x
        for i in range(len(cond_inputs)):
            if cond_inputs[i] is None:
                continue
            label_map = F.interpolate(cond_inputs[i], size=x.size()[2:], mode=self.interpolation)
            if self.separate_projection:
                hidden = self.mlps[i](label_map)
                gamma = self.gammas[i](hidden)
                beta = self.betas[i](hidden)
            else:
                affine_params = self.mlps[i](label_map)
                gamma, beta = affine_params.chunk(2, dim=1)
            if self.bias_only:
                output = output + beta
            else:
                output = output * (1 + gamma) + beta
        return output


class DualAdaptiveNorm(nn.Module):
    def __init__(self,
                 num_features,
                 cond_dims,
                 projection_bias=True,
                 weight_norm_type='',
                 activation_norm_type='instance',
                 activation_norm_params=None,
                 apply_noise=False,
                 bias_only=False,
                 init_gain=1.0,
                 fc_scale=None,
                 is_spatial=None):
        super().__init__()
        if activation_norm_params is None:
            activation_norm_params = SimpleNamespace(affine=False)
        self.mlps = nn.ModuleList()
        self.gammas = nn.ModuleList()
        self.betas = nn.ModuleList()
        self.bias_only = bias_only

        # Make cond_dims a list.
        if type(cond_dims) != list:
            cond_dims = [cond_dims]

        if is_spatial is None:
            is_spatial = [False for _ in range(len(cond_dims))]
        self.is_spatial = is_spatial

        for cond_dim, this_is_spatial in zip(cond_dims, is_spatial):
            kwargs = dict(weight_norm_type=weight_norm_type,
                          bias=projection_bias,
                          init_gain=init_gain,
                          output_scale=fc_scale)
            if this_is_spatial:
                self.gammas.append(Conv2dBlock(cond_dim, num_features, 1, 1, 0, **kwargs))
                self.betas.append(Conv2dBlock(cond_dim, num_features, 1, 1, 0, **kwargs))
            else:
                self.gammas.append(LinearBlock(cond_dim, num_features, **kwargs))
                self.betas.append(LinearBlock(cond_dim, num_features, **kwargs))

        self.norm = get_activation_norm_layer(num_features,
                                              activation_norm_type,
                                              2,
                                              **vars(activation_norm_params))
        self.conditional = True

    def forward(self, x, *cond_inputs, **_kwargs):
        assert len(cond_inputs) == len(self.gammas)
        output = self.norm(x) if self.norm is not None else x
        for cond, gamma_layer, beta_layer in zip(cond_inputs, self.gammas, self.betas):
            if cond is None:
                continue
            gamma = gamma_layer(cond)
            beta = beta_layer(cond)
            if cond.dim() == 4 and gamma.shape != x.shape:
                gamma = F.interpolate(gamma, size=x.size()[2:], mode='bilinear')
                beta = F.interpolate(beta, size=x.size()[2:], mode='bilinear')
            elif cond.dim() == 2:
                gamma = gamma[:, :, None, None]
                beta = beta[:, :, None, None]
            if self.bias_only:
                output = output + beta
            else:
                output = output * (1 + gamma) + beta
        return output


class HyperSpatiallyAdaptiveNorm(nn.Module):
    r"""Spatially Adaptive Normalization (SPADE) initialization.

    Args:
        num_features (int) : Number of channels in the input tensor.
        cond_dims (int or list of int) : List of numbers of channels
            in the conditional input.
        num_filters (int): Number of filters in SPADE.
        kernel_size (int): Kernel size of the convolutional filters in
            the SPADE layer.
        weight_norm_type (str): Type of weight normalization.
            ``'none'``, ``'spectral'``, or ``'weight'``.
        activation_norm_type (str):
            Type of activation normalization.
            ``'none'``, ``'instance'``, ``'batch'``, ``'sync_batch'``,
            ``'layer'``,  ``'layer_2d'``, ``'group'``.
        is_hyper (bool): Whether to use hyper SPADE.
    """

    def __init__(self, num_features, cond_dims,
                 num_filters=0, kernel_size=3,
                 weight_norm_type='',
                 activation_norm_type='sync_batch', is_hyper=True):
        super().__init__()
        padding = kernel_size // 2
        self.mlps = nn.ModuleList()
        if type(cond_dims) != list:
            cond_dims = [cond_dims]

        for i, cond_dim in enumerate(cond_dims):
            mlp = []
            if not is_hyper or (i != 0):
                if num_filters > 0:
                    mlp += [Conv2dBlock(cond_dim, num_filters, kernel_size,
                                        padding=padding,
                                        weight_norm_type=weight_norm_type,
                                        nonlinearity='relu')]
                mlp_ch = cond_dim if num_filters == 0 else num_filters
                mlp += [Conv2dBlock(mlp_ch, num_features * 2, kernel_size,
                                    padding=padding,
                                    weight_norm_type=weight_norm_type)]
                mlp = nn.Sequential(*mlp)
            else:
                if num_filters > 0:
                    raise ValueError('Multi hyper layer not supported yet.')
                mlp = HyperConv2d(padding=padding)
            self.mlps.append(mlp)

        self.norm = get_activation_norm_layer(num_features,
                                              activation_norm_type,
                                              2,
                                              affine=False)

        self.conditional = True

    def forward(self, x, *cond_inputs,
                norm_weights=(None, None), **_kwargs):
        r"""Spatially Adaptive Normalization (SPADE) forward.

        Args:
            x (4D tensor) : Input tensor.
            cond_inputs (list of tensors) : Conditional maps for SPADE.
            norm_weights (5D tensor or list of tensors): conv weights or
            [weights, biases].
        Returns:
            output (4D tensor) : Output tensor.
        """
        output = self.norm(x)
        for i in range(len(cond_inputs)):
            if cond_inputs[i] is None:
                continue
            if type(cond_inputs[i]) == list:
                cond_input, mask = cond_inputs[i]
                mask = F.interpolate(mask, size=x.size()[2:], mode='bilinear', align_corners=False)
            else:
                cond_input = cond_inputs[i]
                mask = None
            label_map = F.interpolate(cond_input, size=x.size()[2:])
            if norm_weights is None or norm_weights[0] is None or i != 0:
                affine_params = self.mlps[i](label_map)
            else:
                affine_params = self.mlps[i](label_map,
                                             conv_weights=norm_weights)
            gamma, beta = affine_params.chunk(2, dim=1)
            if mask is not None:
                gamma = gamma * (1 - mask)
                beta = beta * (1 - mask)
            output = output * (1 + gamma) + beta
        return output


class LayerNorm2d(nn.Module):
    r"""Layer Normalization as introduced in
    https://arxiv.org/abs/1607.06450.
    This is the usual way to apply layer normalization in CNNs.
    Note that unlike the pytorch implementation which applies per-element
    scale and bias, here it applies per-channel scale and bias, similar to
    batch/instance normalization.

    Args:
        num_features (int): Number of channels in the input tensor.
        eps (float, optional, default=1e-5): a value added to the
            denominator for numerical stability.
        affine (bool, optional, default=False): If ``True``, performs
            affine transformation after normalization.
    """

    def __init__(self, num_features, eps=1e-5, channel_only=False, affine=True):
        super(LayerNorm2d, self).__init__()
        self.num_features = num_features
        self.affine = affine
        self.eps = eps
        self.channel_only = channel_only

        if self.affine:
            self.gamma = nn.Parameter(torch.Tensor(num_features).fill_(1.0))
            self.beta = nn.Parameter(torch.zeros(num_features))

    def forward(self, x):
        r"""

        Args:
            x (tensor): Input tensor.
        """
        shape = [-1] + [1] * (x.dim() - 1)
        if self.channel_only:
            mean = x.mean(1, keepdim=True)
            std = x.std(1, keepdim=True)
        else:
            mean = x.view(x.size(0), -1).mean(1).view(*shape)
            std = x.view(x.size(0), -1).std(1).view(*shape)

        x = (x - mean) / (std + self.eps)

        if self.affine:
            shape = [1, -1] + [1] * (x.dim() - 2)
            x = x * self.gamma.view(*shape) + self.beta.view(*shape)
        return x


class ScaleNorm(nn.Module):
    r"""Scale normalization:
    "Transformers without Tears: Improving the Normalization of Self-Attention"
    Modified from:
    https://github.com/tnq177/transformers_without_tears
    """

    def __init__(self, dim=-1, learned_scale=True, eps=1e-5):
        super().__init__()
        # scale = num_features ** 0.5
        if learned_scale:
            self.scale = nn.Parameter(torch.tensor(1.))
        else:
            self.scale = 1.
        # self.num_features = num_features
        self.dim = dim
        self.eps = eps
        self.learned_scale = learned_scale

    def forward(self, x):
        # noinspection PyArgumentList
        scale = self.scale * torch.rsqrt(torch.mean(x ** 2, dim=self.dim, keepdim=True) + self.eps)
        return x * scale

    def extra_repr(self):
        s = 'learned_scale={learned_scale}'
        return s.format(**self.__dict__)


class PixelNorm(ScaleNorm):
    def __init__(self, learned_scale=False, eps=1e-5, **_kwargs):
        super().__init__(1, learned_scale, eps)


class SplitMeanStd(nn.Module):
    def __init__(self, num_features, eps=1e-5, **kwargs):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.multiple_outputs = True

    def forward(self, x):
        b, c, h, w = x.size()
        mean = x.view(b, c, -1).mean(-1)[:, :, None, None]
        var = x.view(b, c, -1).var(-1)[:, :, None, None]
        std = torch.sqrt(var + self.eps)

        # x = (x - mean) / std
        return x, torch.cat((mean, std), dim=1)


class ScaleNorm(nn.Module):
    r"""Scale normalization:
    "Transformers without Tears: Improving the Normalization of Self-Attention"
    Modified from:
    https://github.com/tnq177/transformers_without_tears
    """

    def __init__(self, dim=-1, learned_scale=True, eps=1e-5):
        super().__init__()
        # scale = num_features ** 0.5
        if learned_scale:
            self.scale = nn.Parameter(torch.tensor(1.))
        else:
            self.scale = 1.
        # self.num_features = num_features
        self.dim = dim
        self.eps = eps
        self.learned_scale = learned_scale

    def forward(self, x):
        # noinspection PyArgumentList
        scale = self.scale * torch.rsqrt(
            torch.mean(x ** 2, dim=self.dim, keepdim=True) + self.eps)
        return x * scale

    def extra_repr(self):
        s = 'learned_scale={learned_scale}'
        return s.format(**self.__dict__)


class PixelLayerNorm(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.norm = nn.LayerNorm(*args, **kwargs)

    def forward(self, x):
        if x.dim() == 4:
            b, c, h, w = x.shape
            return self.norm(x.permute(0, 2, 3, 1).view(-1, c)).view(b, h, w, c).permute(0, 3, 1, 2)
        else:
            return self.norm(x)


def get_activation_norm_layer(num_features, norm_type, input_dim, **norm_params):
    r"""Return an activation normalization layer.

    Args:
        num_features (int): Number of feature channels.
        norm_type (str):
            Type of activation normalization.
            ``'none'``, ``'instance'``, ``'batch'``, ``'sync_batch'``,
            ``'layer'``,  ``'layer_2d'``, ``'group'``, ``'adaptive'``,
            ``'spatially_adaptive'`` or ``'hyper_spatially_adaptive'``.
        input_dim (int): Number of input dimensions.
        norm_params: Arbitrary keyword arguments that will be used to
            initialize the activation normalization.
    """
    input_dim = max(input_dim, 1)  # Norm1d works with both 0d and 1d inputs

    if norm_type == 'none' or norm_type == '':
        norm_layer = None
    elif norm_type == 'batch':
        norm = getattr(nn, 'BatchNorm%dd' % input_dim)
        norm_layer = norm(num_features, **norm_params)
    elif norm_type == 'instance':
        affine = norm_params.pop('affine', True)  # Use affine=True by default
        norm = getattr(nn, 'InstanceNorm%dd' % input_dim)
        norm_layer = norm(num_features, affine=affine, **norm_params)
    elif norm_type == 'sync_batch':
        norm_layer = SyncBatchNorm(num_features, **norm_params)
    elif norm_type == 'layer':
        norm_layer = nn.LayerNorm(num_features, **norm_params)
    elif norm_type == 'layer_2d':
        norm_layer = LayerNorm2d(num_features, **norm_params)
    elif norm_type == 'pixel_layer':
        elementwise_affine = norm_params.pop('affine', True)  # Use affine=True by default
        norm_layer = PixelLayerNorm(num_features, elementwise_affine=elementwise_affine, **norm_params)
    elif norm_type == 'scale':
        norm_layer = ScaleNorm(**norm_params)
    elif norm_type == 'pixel':
        norm_layer = PixelNorm(**norm_params)
        import imaginaire.config
        if imaginaire.config.USE_JIT:
            norm_layer = torch.jit.script(norm_layer)
    elif norm_type == 'group':
        num_groups = norm_params.pop('num_groups', 4)
        norm_layer = nn.GroupNorm(num_channels=num_features, num_groups=num_groups, **norm_params)
    elif norm_type == 'adaptive':
        norm_layer = AdaptiveNorm(num_features, **norm_params)
    elif norm_type == 'dual_adaptive':
        norm_layer = DualAdaptiveNorm(num_features, **norm_params)
    elif norm_type == 'spatially_adaptive':
        if input_dim != 2:
            raise ValueError('Spatially adaptive normalization layers '
                             'only supports 2D input')
        norm_layer = SpatiallyAdaptiveNorm(num_features, **norm_params)
    elif norm_type == 'hyper_spatially_adaptive':
        if input_dim != 2:
            raise ValueError('Spatially adaptive normalization layers '
                             'only supports 2D input')
        norm_layer = HyperSpatiallyAdaptiveNorm(num_features, **norm_params)
    else:
        raise ValueError('Activation norm layer %s '
                         'is not recognized' % norm_type)
    return norm_layer