models.py

import torch
from torch import nn
import torch.nn.functional as F
from torch import Tensor
from typing import Optional, List
from timm.models.layers import trunc_normal_


class LayerNorm(nn.Module):
    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 
    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
    shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
    with shape (batch_size, channels, height, width).
    """
    def __init__(self, normalized_shape, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.eps = eps
        self.normalized_shape = (normalized_shape, )
    
    def forward(self, x):
        u = x.mean(1, keepdim=True)
        s = (x - u).pow(2).mean(1, keepdim=True)
        x = (x - u) / torch.sqrt(s + self.eps)
        x = self.weight[:, None, None] * x + self.bias[:, None, None]
        return x


def conv_block(
    din: int,
    dout: int,
    k: int,
    act: nn.Module = nn.Identity,
    depthwise: bool = False,
    bottleneck: bool = False,
    batchnorm: bool = False,
    separable: bool = False,
    batchnorm_type: str = "bn",
    **kwargs
):
    if batchnorm_type == "bn":
        bn_fun = lambda x: nn.BatchNorm2d(x, eps=1e-2)
    elif batchnorm_type == "frn":
        bn_fun = lambda x: FRN(x)
    else:
        raise NotImplementedError(f"batchnorm_type {batchnorm_type} not implemented")

    if not depthwise:
        d1 = din // 4 if bottleneck else din
        kwargs = kwargs.copy()
        kwargs["bias"] = True # (not batchnorm)
        batchnorm = True # (not batchnorm)
        mods = []
        mods.append(nn.Conv2d(d1, dout, k, **kwargs))
        if batchnorm:
            mods.append(bn_fun(dout))
        mods.append(act())
        if bottleneck:
            blayer = [nn.Conv2d(din, din // 4, 1, bias=(not batchnorm))]
            if batchnorm:
                blayer.append(LayerNorm(din // 4, eps=1e-2))
            mods = mods + blayer
    elif separable:
        kwargs = kwargs.copy()
        kwargs2 = kwargs.copy()
        kwargs3 = kwargs.copy()
        kwargs["groups"] = din
        kwargs["padding"] = [0, k//2]
        kwargs["bias"] = True
        kwargs2["groups"] = din
        kwargs2["padding"] = [k//2, 0]
        kwargs2["bias"] = False
        kwargs3["groups"] = 1
        kwargs3["padding"] = 0
        kwargs3["bias"] = (not batchnorm)
        mods = [
            nn.Conv2d(din, din, (1, k), **kwargs),
            nn.Conv2d(din, din, (k, 1), **kwargs2),
            nn.Conv2d(din, dout, 1, **kwargs3),
        ]
        if batchnorm:
            mods.append(bn_fun(dout))
        mods.append(act())
    else:
        kwargs = kwargs.copy()
        kwargs2 = kwargs.copy()
        kwargs["groups"] = din
        kwargs["padding"] = k//2
        kwargs["bias"] = True
        kwargs2["groups"] = 1
        kwargs2["padding"] = 0
        kwargs2["bias"] = True  # 
        mods = [
            nn.Conv2d(din, din, (k, k), **kwargs),
            nn.Conv2d(din, dout, 1, **kwargs2)
        ]
        if batchnorm:
            mods.append(bn_fun(dout))
        mods.append(act())
    mod = nn.Sequential(*mods)
    return mod


class FRN(nn.Module):
    def __init__(self, num_features, eps=1e-6, is_eps_leanable=False):
        """
        weight = gamma, bias = beta
        beta, gamma:
            Variables of shape [1, 1, 1, C]. if TensorFlow
            Variables of shape [1, C, 1, 1]. if PyTorch
        eps: A scalar constant or learnable variable.
        """
        super(FRN, self).__init__()

        self.num_features = num_features
        self.init_eps = eps
        self.is_eps_leanable = is_eps_leanable

        self.weight = nn.parameter.Parameter(
            torch.Tensor(1, num_features, 1, 1), requires_grad=True)
        self.bias = nn.parameter.Parameter(
            torch.Tensor(1, num_features, 1, 1), requires_grad=True)
        if is_eps_leanable:
            self.eps = nn.parameter.Parameter(torch.Tensor(1), requires_grad=True)
        else:
            self.register_buffer('eps', torch.Tensor([eps]))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.ones_(self.weight)
        nn.init.zeros_(self.bias)
        if self.is_eps_leanable:
            nn.init.constant_(self.eps, self.init_eps)

    def extra_repr(self):
        return 'num_features={num_features}, eps={init_eps}'.format(**self.__dict__)

    def forward(self, x):
        """
        0, 1, 2, 3 -> (B, H, W, C) in TensorFlow
        0, 1, 2, 3 -> (B, C, H, W) in PyTorch
        TensorFlow code
            nu2 = tf.reduce_mean(tf.square(x), axis=[1, 2], keepdims=True)
            x = x * tf.rsqrt(nu2 + tf.abs(eps))
            # This Code include TLU function max(y, tau)
            return tf.maximum(gamma * x + beta, tau)
        """
        # Compute the mean norm of activations per channel.
        nu2 = x.pow(2).mean(dim=[2, 3], keepdim=True)

        # Perform FRN.
        x = x * torch.rsqrt(nu2 + self.eps.abs())

        # Scale and Bias
        x = self.weight * x + self.bias
        return x



class res_layers(nn.Module):
    def __init__(self, d: int, n: int, k: int, act: nn.Module, **kwargs):
        super().__init__()
        kwargs = kwargs.copy()
        self.convs = nn.ModuleList(
            [conv_block(d, d, k, padding=k//2, batchnorm=True, act=nn.Identity) for _ in range(n)]
        )
        self.acts = nn.ModuleList([act() for _ in range(n)])

    def forward(self, inputs: Tensor):
        x = inputs
        for f, act in zip(self.convs, self.acts):
            x = act(f(x) + x)
        return x


def down_layer(d: int, k: int, act: nn.Module, num_res=0, pool=nn.MaxPool2d, **kwargs):
    modules = [
        pool(2, 2, 0), conv_block(d // 2, d, k, act, **kwargs)
    ]
    if num_res > 0:
        modules.append(res_layers(d, num_res, k, act, **kwargs))
    return nn.Sequential(*modules)


class up_layer(nn.Module):
    def __init__(self, d: int, k: int, act: nn.Module, num_res=0, **kwargs):
        super().__init__()
        modules = []
        if num_res > 0:
            modules.append(res_layers(2 * d, num_res, k, act, **kwargs))
        modules.append(nn.UpsamplingBilinear2d(scale_factor=2))
        modules.append(conv_block(2 * d, d, k, act, **kwargs))
        self.net1 = nn.Sequential(*modules)
        self.net2 = conv_block(2 * d, d, k, act, **kwargs)

    def forward(self, x: Tensor, z: Tensor):
        x = self.net1(x)
        x = torch.cat([x, z], 1)
        x = self.net2(x)
        return x


class UNetEncoder(nn.Module):
    def __init__(
        self,
        cin,
        cout,
        ksize=3,
        depth=3,
        num_res=1,
        groups: int = 1,
        n_hidden=16,
        act=nn.SiLU,
        pool='MaxPool2d',
        dropout: bool = False,
        bottleneck: bool = False,
        depthwise: bool = False,
        separable: bool = False,
        batchnorm: bool = False,
        batchnorm_type: str = "bn",
    ) -> None:
        super().__init__()
        self.cin = cin
        self.cout = cout
        self.num_res = num_res
        self.depth = depth
        self.n_hidden = n_hidden
        self.ksize = ksize
        self.groups = groups
        d, k = n_hidden, ksize
        self.act = act
        self.batchnorm = batchnorm
        act = getattr(nn, act) if isinstance(act, str) else act
        pool = getattr(nn, pool)

        kwargs = dict(
            bias=True,
            groups=groups,
            padding=k // 2,
            padding_mode="replicate",
            bottleneck=bottleneck,
            depthwise=depthwise,
            separable=separable,
            batchnorm=batchnorm,
            batchnorm_type=batchnorm_type,
        )

        kwargs1 = kwargs.copy()
        kwargs1["groups"] = 1
        self.first = nn.Sequential(
            nn.Conv2d(cin, d, k, padding=k//2),
            # LayerNorm(d, eps=1e-6),
            FRN(d, eps=1e-2),
            act(),
            # conv_block(cin, d, k, act, **kwargs1),
            res_layers(d, num_res, k, act, **kwargs)
        )

        self.down = nn.ModuleList()
        for _ in range(depth):
            d *= 2
            layer = down_layer(d, k, act, num_res=num_res, pool=pool, **kwargs)
            self.down.append(layer)

        if dropout:
            self.dropout = nn.Dropout2d(0.5)

        self.up = nn.ModuleList()
        for _ in range(depth):
            d //= 2
            layer = up_layer(d, k, act, num_res=num_res, **kwargs)
            self.up.append(layer)

        kwargsf = kwargs.copy()
        kwargsf["groups"] = 1
        self.final = nn.Sequential(
            res_layers(d, num_res, k, act, **kwargs),
            nn.Conv2d(d, cout, k, padding=k//2)
            # conv_block(d, cout, k, **kwargsf)
        )
        self.apply(self._init_weights)

    def forward(self, inputs: Tensor):
        x = inputs
        x = self.first(x)

        # downsampling
        intermediate_outputs = [x]
        for _, fdown in enumerate(self.down):
            x = fdown(x)
            intermediate_outputs.append(x)

        # dropout
        if hasattr(self, "dropout"):
            x = self.dropout(x)

        # upsampling
        for i, fup in enumerate(self.up):
            z = intermediate_outputs[-(i + 2)]
            x = fup(x, z)

        # final layers and desired size
        x = self.final(x)

        return x
    
    def _init_weights(self, m):
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            trunc_normal_(m.weight, std=.02)
            nn.init.constant_(m.bias, 0)


class Decoder(nn.Module):
    def __init__(self, din, dout, n_hidden=64, n_res=0, act=nn.SiLU, offset=False, collapse_dims=True, loss_type="mse", fit_sigma: bool = False, **kwargs):
        super().__init__()
        if offset:
            self.spatial_emb = SpatialEmbeddings(din, n_hidden, n_hidden)
        self.din = din
        self.dout = dout
        self.collapse_dims =collapse_dims
        self.hin = n_hidden if offset else din
        self.head = nn.Sequential(
            conv_block(self.hin, n_hidden, 1, act, **kwargs),
            *[conv_block(n_hidden, n_hidden, 1, act, **kwargs) for _ in range(n_res)],
            nn.Conv2d(n_hidden, dout, 1),
        )
        self.loss_type = loss_type
        if loss_type == "mse":
            self.log_sig = nn.Parameter(torch.tensor(0.0), requires_grad=fit_sigma)

        # self.head = conv_block(n_hidden, n_hidden, 1, act, **kwargs)

    def forward_offset(self, latent, offset):
        spatial = self.spatial_emb(offset)
        bs, D, D1 = spatial.shape
        *_, nr, nc = latent.shape
        spatial = spatial.view(bs, D, D1, 1, 1)
        latent = latent.view(bs, D, 1, nr, nc)
        v = (latent * spatial).view(bs, D, D1, nr, nc)
        out = v.sum(1) # .view(1, D, nr, nc)
        out = self.head(out)
        return out

    def forward(self, latent, offset=None):
        if offset is not None:
            out = self.forward_offset(latent, offset)
        else:
            out = self.head(latent)
        if self.collapse_dims:
            out = out.squeeze(1)
        return out

    def loss_binary(
        self,
        tgt: Tensor,
        C: Optional[Tensor] = None,
        M: Optional[Tensor] = None,
    ):
        if M is None:
            M = torch.ones_like(tgt)
        L = self(C)
        if len(L.shape) == 0:
            L = torch.ones_like(tgt) * L
        # L = L.clip(-5.0, 5.0)
        # loss = torch.where(
        #     L >= 0.0,
        #     torch.log(1.0 + torch.exp(-L)) + (1.0 - tgt) * L,
        #     torch.log(1.0 + torch.exp(L)) - tgt * L,
        # )
        loss = F.binary_cross_entropy_with_logits(L, tgt, reduction='none')
        loss = (loss * M).sum() / tgt.shape[0]
        return L, loss
    
    def loss_mse(
        self,
        Y: Tensor,
        C: Optional[Tensor] = None,
        A: Optional[Tensor] = None,
        M: Optional[Tensor] = None,
    ):
        if M is None:
            M = torch.ones_like(Y)
        Yhat = self.forward(C, A) * M
        # return Yhat, 0.0
        sig = (1e-4 + self.log_sig).exp()
        prec = 1.0 / (1e-4 + sig ** 2)
        err = 0.5 * prec * F.mse_loss(Yhat, Y, reduction="none") + self.log_sig
        prior = sig - self.log_sig
        # err = 0.5 * F.mse_loss(Yhat, Y, reduction="none") * M
        loss = ((err * M).sum() + prior)
        # sigerr = F.smooth_l1_loss(mse.detach().m,sqrt(), sig)
        loss = loss  / Y.shape[0]
        return Yhat, loss

    def loss(self, *args, **kwargs):
        if self.loss_type == "binary":
            l = self.loss_binary(*args, **kwargs)
        elif self.loss_type == "mse":
            l = self.loss_mse(*args, **kwargs)
        return l 


class SpatialEmbeddings(nn.Module):
    def __init__(self, dout, nl, n_hidden=64, act=nn.SiLU, n_chunks=1):
        super().__init__()
        self.D = dout
        self.nl = nl
        act = getattr(nn, act) if isinstance(act, str) else act
        self.dec = nn.Sequential(
            nn.Linear(2, n_hidden),
            act(),
            nn.Linear(n_hidden, n_hidden),
            act(),
            nn.Linear(n_hidden, dout * nl),
        )

    def forward(self, inputs):
        u = inputs
        u = (
            [u] #+
            # [torch.sin(u / (1000 ** (2 * k / 12))) for k in range(0, 11, 2)]
            # + [torch.cos(u / (1000 ** (2 * k / 12))) for k in range(0, 11, 2)]
            # [torch.sin(2 * np.pi * u / k) for k in range(1, 12, 2)] +
            # [torch.cos(2 * np.pi * u / k) for k in range(1, 12, 2)]
        )
        u = torch.cat(u, -1)
        out = self.dec(u)
        out = out.view(-1, self.D, self.nl)
        return out