DeepTime/models/modules/encoders.py

# import gin
import torch
import torch.nn as nn
from torch import Tensor
from einops import rearrange, repeat, reduce

from tsai.models.InceptionTimePlus import (
    Conv,
    noop,
    nn,
    LinBnDrop,
    GAP1d,
    torch,
    AddCoords1d, BatchNorm
)
from tsai.models.TSTPlus import TSTPlus
from tsai.models.TSPerceiver import TSPerceiver
from tsai.models.TSSequencerPlus import TSSequencerPlus
from torch.nn.utils import weight_norm, spectral_norm
from .causalinception import CausalInceptionTimePlus
from .inr import INR

def custom_head(head_nf, c_out, seq_len):
    return nn.Sequential(
        # AddCoords1d(),
        # Conv(head_nf+1, head_nf, 2, bias=True, norm='Spectral'),
        # nn.BatchNorm1d(head_nf),
        # # nn.Dropout(0.15),
        # nn.ReLU(),
        AddCoords1d(),
        Conv(head_nf + 1, c_out, 1, bias=False, norm="Spectral"),
    )
    
class LinBnDropSN(nn.Sequential):
    "Module grouping `BatchNorm1d`, `Dropout` and `Linear` layers"
    def __init__(self, n_in, n_out, bn=True, p=0., act=None, lin_first=False, norm=None):
        layers = [BatchNorm(n_out if lin_first else n_in, ndim=1)] if bn else []
        if p != 0: layers.append(nn.Dropout(p))
        lin = [spectral_norm(nn.Linear(n_in, n_out, bias=not bn))]
        if act is not None: lin.append(act)
        layers = lin+layers if lin_first else layers+lin
        super().__init__(*layers)


class InceptionEncoder(nn.Module):
    def __init__(self, c_in, c_out, dropout, layers, layer_size, *args, **kwargs):
        super().__init__()
        self.net = CausalInceptionTimePlus(
            c_in=c_in, c_out=c_out, ks=[39, 19, 3], custom_head=custom_head, coord=True, fc_dropout=dropout, bn=True, depth=layers, nf=layer_size, *args, **kwargs
        )
        bn = kwargs.get("bn", True)
        fc_dropout = kwargs.get("fc_dropout", 0.15)
        self.pool = nn.Sequential(
            # GACP1d(1),
            # LinBnDrop(c_out*2, c_out, bn=bn, p=dropout)
            GAP1d(1),
            LinBnDropSN(c_out, c_out, bn=bn, p=fc_dropout),
        )
        self.head = nn.Sequential(
            # just to make sure we get a spectral norm final layer (after cat)
            LinBnDropSN(c_out*2, c_out, bn=bn, p=fc_dropout),
        )

    def forward(self, x):
        """
        Takes in a sequence of shape (batch, sequence, features)
        and outputs a representation of shape (batch, features)
        """
        outs = self.net(x.permute(0, 2, 1))  # .permute(0, 2, 1)
        last = outs[:, :, -1]  # take last
        max = self.pool(outs)
        return self.head(torch.cat([max, last], 1))


class TransformerEncoder(nn.Module):
    def __init__(
        self,
        c_in,
        c_out,
        seq_len,
        layers=3,
        layer_size=512,
        dropout=0.1,
        n_heads=4,
        conv_dropout=0,
        *args,
        **kwargs,
    ):
        super().__init__()
        # d_model (82) must be divisible by n_heads (4)
        layer_size = layer_size // n_heads * n_heads
        d_model = layer_size // 2
        self.net = TSTPlus(
            c_in=c_in,
            c_out=c_out,
            seq_len=seq_len,
            d_model=d_model,
            n_heads=n_heads,
            d_k=d_model // n_heads,
            d_v=d_model // n_heads,
            d_ff=layer_size,
            n_layers=layers,
            dropout=conv_dropout,
            fc_dropout=dropout,
            flatten=False,
            # *args, **kwargs
        )

    def forward(self, x):
        """
        Takes in a sequence of shape (batch, sequence, features)
        and outputs a representation of shape (batch, features)
        """
        outs = self.net(x.permute(0, 2, 1))
        return outs
    

class TransformerEncoder2(nn.Module):
    def __init__(
        self,
        c_in,
        c_out,
        seq_len,
        layers=3,
        layer_size=512,
        dropout=0.1,
        n_heads=4,
        conv_dropout=0,
        *args,
        **kwargs,
    ):
        super().__init__()
        # d_model (82) must be divisible by n_heads (4)
        layer_size = layer_size // n_heads * n_heads
        d_model = layer_size // 2 
        self.net = TSPerceiver(
            c_in=c_in,
            c_out=c_out,
            seq_len=seq_len,
            
            # cat_szs=0, n_cont=0, 
            n_latents=layer_size, d_latent=layer_size//4, 
            # d_context=None, 
            self_per_cross_attn=1,
                #  share_weights=True, cross_n_heads=1, d_head=None, 
                 
                 
            # d_model=d_model,
            # d_k=d_model // n_heads,
            # d_v=d_model // n_heads,
            # d_ff=layer_size,
            
            self_n_heads=n_heads,
            n_layers=layers,
            attn_dropout=conv_dropout,
            fc_dropout=dropout,
        )

    def forward(self, x):
        """
        Takes in a sequence of shape (batch, sequence, features)
        and outputs a representation of shape (batch, features)
        """
        outs = self.net(x.permute(0, 2, 1))
        return outs


class LSTMEncoder(nn.Module):
    def __init__(
        self,
        c_in,
        c_out,
        dropout=0.1,
        conv_dropout=0,
        layers=1,
        layer_size=100,
        *args,
        **kwargs,
    ):
        super().__init__()
        self.rnn = nn.LSTM(
            c_in,
            layer_size,
            num_layers=layers,
            bias=True,
            batch_first=True,
            dropout=conv_dropout,
            bidirectional=False,
        )
        self.dropout = nn.Dropout(dropout) if dropout else noop
        self.fc = nn.Linear(layer_size * (1 + 0), c_out)

    def forward(self, x):
        """
        Takes in a sequence of shape (batch, sequence, features)
        and outputs a representation of shape (batch, features)
        """
        # x = x.transpose(2,1)    # [batch_size x n_vars x seq_len] --> [batch_size x seq_len x n_vars]
        output, _ = self.rnn(
            x
        )  # output from all sequence steps: [batch_size x seq_len x hidden_size * (1 + bidirectional)]
        output = output[
            :, -1
        ]  # output from last sequence step : [batch_size x hidden_size * (1 + bidirectional)]
        return self.fc(self.dropout(output))



class LSTMEncoder2(nn.Module):
    def __init__(
        self,
        c_in,
        c_out,
        seq_len,
        dropout=0.1,
        conv_dropout=0,
        layers=1,
        layer_size=100,
        *args,
        **kwargs,
    ):
        super().__init__()
        self.rnn = TSSequencerPlus(
            c_in=c_in,
            c_out=c_out,
            seq_len=seq_len,
            d_model=layer_size,
            depth=layers,
            lstm_dropout=conv_dropout,
            fc_dropout=dropout,
            pre_norm=False, use_token=True, use_pe=True, 
            use_bn=False, 
        )

    def forward(self, x):
        """
        Takes in a sequence of shape (batch, sequence, features)
        and outputs a representation of shape (batch, features)
        """
        return self.rnn(
            x.transpose(2, 1)
        )


class MLPEncoder(nn.Module):
    def __init__(
        self,
        c_in,
        c_out,
        scales=[0.01, 0.1, 1, 5, 10, 20, 50, 100],
        n_fourier_feats=4096,
        layers=2,
        layer_size=32,
        *args,
        **kwargs,
    ):
        super().__init__()
        self.net = INR(
            in_feats=c_in,
            out_feats=layer_size,
            scales=scales,
            n_fourier_feats=n_fourier_feats,
            layers=layers,
            layer_size=layer_size,
        )
        self.head = nn.Linear(layer_size, c_out)

    def forward(self, x):
        """
        Takes in a sequence of shape (batch, sequence, features)
        and outputs a representation of shape (batch, features)
        """
        return self.head(self.net(x)[:, -1])