[ENH] ModernTCN: A Modern Pure Convolution Structure for General Time Series Analysis added

docs/source/m_layer_v2.rst

@@ -53,3 +53,4 @@ See the detailed API documentation for the V2 base classes and specific model im
    models.scinet._scinet_v2.SCINet_v2
    models.patch_tst._patch_tst_v2.PatchTST_v2
    models.frets._frets_v2.FreTS
+   models.modern_tcn._modern_tcn_v2.ModernTCN

docs/source/pkg_v2.rst

@@ -105,3 +105,4 @@ See the detailed API documentation for the available V2 Package classes below:
    models.scinet._scinet_pkg_v2.SCINet_pkg_v2
    models.patch_tst._patch_tst_pkg_v2.PatchTST_pkg_v2
    models.frets._frets_pkg_v2.FreTS_pkg_v2
+   models.modern_tcn._modern_tcn_pkg_v2.ModernTCN_pkg_v2

pytorch_forecasting/layers/_blocks/__init__.py

@@ -1,4 +1,5 @@
 from pytorch_forecasting.layers._blocks._frets_block import FreTSCore
+from pytorch_forecasting.layers._blocks._modern_tcn_block import ModernTCNBlock
 from pytorch_forecasting.layers._blocks._residual_block_dsipts import ResidualBlock
 from pytorch_forecasting.layers._blocks._scinet_block import SCIBlock
 from pytorch_forecasting.layers._blocks._softs_block import (
@@ -8,6 +9,7 @@
 __all__ = [
     "FreTSCore",
     "ResidualBlock",
+    "ModernTCNBlock",
     "SCIBlock",
     "STADModule",
 ]

pytorch_forecasting/layers/_blocks/_modern_tcn_block.py

@@ -0,0 +1,95 @@
+"""
+ModernTCN Block: For Modern Temporal Convolutional Network
+"""
+
+import torch
+import torch.nn as nn
+
+from pytorch_forecasting.layers._convolution._reparam_large_kernel_conv import (
+    ReparamLargeKernelConv,
+)
+
+
+class ModernTCNBlock(nn.Module):
+    """
+    Modern TCN Block.
+
+    This block is a residual block that consists of a depthwise
+    separable convolution and a feed-forward network.
+
+    Parameters
+    ----------
+    d_model : int
+        Dimension of the model.
+    kernel_size : int
+        Size of the large kernel.
+    small_kernel_size : int
+        Size of the small kernel.
+    d_ff : int
+        Dimension of the feed-forward network.
+    nvars : int
+        Number of variables.
+    dropout : float
+        Dropout rate.
+    """
+
+    def __init__(self, d_model, kernel_size, small_kernel_size, d_ff, nvars, dropout):
+        super().__init__()
+        self.d_model = d_model
+        self.nvars = nvars
+
+        self.dwconv = ReparamLargeKernelConv(
+            in_channels=nvars * d_model,
+            out_channels=nvars * d_model,
+            kernel_size=kernel_size,
+            stride=1,
+            groups=nvars * d_model,
+            small_kernel_size=small_kernel_size,
+        )
+        self.norm = nn.BatchNorm1d(d_model)
+
+        self.ffn1_pw1 = nn.Conv1d(
+            nvars * d_model, nvars * d_ff, kernel_size=1, groups=nvars
+        )
+        self.ffn1_act = nn.GELU()
+        self.ffn1_pw2 = nn.Conv1d(
+            nvars * d_ff, nvars * d_model, kernel_size=1, groups=nvars
+        )
+        self.ffn1_drop1 = nn.Dropout(dropout)
+        self.ffn1_drop2 = nn.Dropout(dropout)
+
+        self.ffn2_pw1 = nn.Conv1d(
+            nvars * d_model, nvars * d_ff, kernel_size=1, groups=d_model
+        )
+        self.ffn2_act = nn.GELU()
+        self.ffn2_pw2 = nn.Conv1d(
+            nvars * d_ff, nvars * d_model, kernel_size=1, groups=d_model
+        )
+        self.ffn2_drop1 = nn.Dropout(dropout)
+        self.ffn2_drop2 = nn.Dropout(dropout)
+
+    def forward(self, x):
+        input_x = x
+        B, M, D, N = x.shape
+
+        x = x.reshape(B, M * D, N)
+        x = self.dwconv(x)
+
+        x = x.reshape(B * M, D, N)
+        x = self.norm(x)
+        x = x.reshape(B, M * D, N)
+
+        x = self.ffn1_drop1(self.ffn1_pw1(x))
+        x = self.ffn1_act(x)
+        x = self.ffn1_drop2(self.ffn1_pw2(x))
+        x = x.reshape(B, M, D, N)
+
+        x = x.permute(0, 2, 1, 3)
+        x = x.reshape(B, D * M, N)
+        x = self.ffn2_drop1(self.ffn2_pw1(x))
+        x = self.ffn2_act(x)
+        x = self.ffn2_drop2(self.ffn2_pw2(x))
+        x = x.reshape(B, D, M, N)
+        x = x.permute(0, 2, 1, 3)
+
+        return input_x + x

pytorch_forecasting/layers/_convolution/__init__.py

@@ -0,0 +1,5 @@
+from pytorch_forecasting.layers._convolution._reparam_large_kernel_conv import (
+    ReparamLargeKernelConv,
+)
+
+__all__ = ["ReparamLargeKernelConv"]

pytorch_forecasting/layers/_convolution/_reparam_large_kernel_conv.py

@@ -0,0 +1,67 @@
+"""
+Reparameterizable Large Kernel Convolution.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class ReparamLargeKernelConv(nn.Module):
+    """
+    Reparameterizable Large Kernel Convolution.
+
+    This layer uses a large kernel (kernel_size) and
+    a small kernel in parallel,then adds their outputs.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels.
+    out_channels : int
+        Number of output channels.
+    kernel_size : int
+        Large kernel size.
+    stride : int
+        Stride.
+    groups : int
+        Number of groups.
+    small_kernel_size : int
+        Small kernel size.
+    """
+
+    def __init__(
+        self, in_channels, out_channels, kernel_size, stride, groups, small_kernel_size
+    ):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.small_kernel_size = small_kernel_size
+
+        padding = kernel_size // 2
+        self.lkb_origin = nn.Sequential(
+            nn.Conv1d(
+                in_channels,
+                out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+                groups=groups,
+                bias=False,
+            ),
+            nn.BatchNorm1d(out_channels),
+        )
+
+        self.small_conv = nn.Sequential(
+            nn.Conv1d(
+                in_channels,
+                out_channels,
+                kernel_size=small_kernel_size,
+                stride=stride,
+                padding=small_kernel_size // 2,
+                groups=groups,
+                bias=False,
+            ),
+            nn.BatchNorm1d(out_channels),
+        )
+
+    def forward(self, x):
+        return self.lkb_origin(x) + self.small_conv(x)

pytorch_forecasting/layers/_head/__init__.py

@@ -0,0 +1,3 @@
+from pytorch_forecasting.layers._head._flatten_head import Flatten_Head
+
+__all__ = ["Flatten_Head"]

pytorch_forecasting/layers/_head/_flatten_head.py

@@ -0,0 +1,61 @@
+"""
+Flatten Head layer for time series forecasting models.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class Flatten_Head(nn.Module):
+    """
+    Flatten Head.
+
+    This layer flattens the input and projects
+    it to the target window.
+
+    Parameters
+    ----------
+    individual : bool
+        If True, uses a separate linear projection per variable.
+    n_vars : int
+        Number of variables.
+    nf : int
+        Number of features.
+    target_window : int
+        Length of the target window.
+    head_dropout : float
+        Dropout rate.
+    """
+
+    def __init__(self, individual, n_vars, nf, target_window, head_dropout=0):
+        super().__init__()
+        self.individual = individual
+        self.n_vars = n_vars
+
+        if self.individual:
+            self.linears = nn.ModuleList()
+            self.dropouts = nn.ModuleList()
+            self.flattens = nn.ModuleList()
+            for _ in range(self.n_vars):
+                self.flattens.append(nn.Flatten(start_dim=-2))
+                self.linears.append(nn.Linear(nf, target_window))
+                self.dropouts.append(nn.Dropout(head_dropout))
+        else:
+            self.flatten = nn.Flatten(start_dim=-2)
+            self.linear = nn.Linear(nf, target_window)
+            self.dropout = nn.Dropout(head_dropout)
+
+    def forward(self, x):
+        if self.individual:
+            x_out = []
+            for i in range(self.n_vars):
+                z = self.flattens[i](x[:, i, :, :])
+                z = self.linears[i](z)
+                z = self.dropouts[i](z)
+                x_out.append(z)
+            x = torch.stack(x_out, dim=1)
+        else:
+            x = self.flatten(x)
+            x = self.linear(x)
+            x = self.dropout(x)
+        return x

pytorch_forecasting/models/__init__.py

@@ -12,6 +12,7 @@
 from pytorch_forecasting.models.deepar import DeepAR
 from pytorch_forecasting.models.frets import FreTS, FreTS_pkg_v2
 from pytorch_forecasting.models.mlp import DecoderMLP
+from pytorch_forecasting.models.modern_tcn import ModernTCN
 from pytorch_forecasting.models.nbeats import NBeats, NBeatsKAN
 from pytorch_forecasting.models.nhits import NHiTS
 from pytorch_forecasting.models.nn import GRU, LSTM, MultiEmbedding, get_rnn
@@ -59,6 +60,7 @@
     "SOFTS_pkg_v2",
     "SCINet_v2",
     "SCINet_pkg_v2",
+    "ModernTCN",
     "FreTS",
     "FreTS_pkg_v2",
 ]

pytorch_forecasting/models/modern_tcn/__init__.py

@@ -0,0 +1,6 @@
+"""Modern Temporal Convolutional Network model for time series forecasting."""
+
+from pytorch_forecasting.models.modern_tcn._modern_tcn_pkg_v2 import ModernTCN_pkg_v2
+from pytorch_forecasting.models.modern_tcn._modern_tcn_v2 import ModernTCN
+
+__all__ = ["ModernTCN", "ModernTCN_pkg_v2"]

pytorch_forecasting/models/modern_tcn/_modern_tcn_pkg_v2.py

@@ -0,0 +1,93 @@
+"""ModernTCN package container."""
+
+from pytorch_forecasting.base._base_pkg import Base_pkg
+
+
+class ModernTCN_pkg_v2(Base_pkg):
+    """
+    ModernTCN package container
+
+    GitHub Repository:https://github.com/luodhhh/ModernTCN
+
+    Research Paper: https://openreview.net/forum?id=vpJMJerXHU
+
+    """
+
+    _tags = {
+        "info:name": "ModernTCN",
+        "authors": ["Muhammad-Rebaal", "luodhhh"],
+        "info:compute": 2,
+        "info:y_type": ["numeric"],
+        "capability:exogenous": True,
+        "capability:multivariate": True,
+        "capability:pred_int": True,
+        "capability:flexible_history_length": False,
+        "capability:cold_start": False,
+    }
+
+    @classmethod
+    def get_cls(cls):
+        """Get model class."""
+        from pytorch_forecasting.models.modern_tcn._modern_tcn_v2 import ModernTCN
+
+        return ModernTCN
+
+    @classmethod
+    def get_datamodule_cls(cls):
+        """Get the underlying DataModule class."""
+        from pytorch_forecasting.data.data_module import (
+            EncoderDecoderTimeSeriesDataModule,
+        )
+
+        return EncoderDecoderTimeSeriesDataModule
+
+    @classmethod
+    def get_test_train_params(cls):
+        """Return testing parameter settings."""
+        from pytorch_forecasting.metrics import QuantileLoss
+
+        params = [
+            {},
+            {
+                "d_model": 16,
+                "kernel_size": 3,
+                "n_blocks": 1,
+                "d_ff": 32,
+                "patch_size": 4,
+                "use_revin": False,
+            },
+            {
+                "d_model": 8,
+                "kernel_size": 3,
+                "n_blocks": 1,
+                "d_ff": 16,
+                "patch_size": 2,
+                "use_revin": True,
+            },
+            {
+                "loss": QuantileLoss(quantiles=[0.1, 0.5, 0.9]),
+                "d_model": 16,
+                "kernel_size": 3,
+                "n_blocks": 1,
+                "d_ff": 32,
+                "patch_size": 4,
+                "use_revin": False,
+            },
+            {
+                "d_model": 16,
+                "kernel_size": 7,
+                "small_kernel_size": 3,
+                "n_blocks": 1,
+                "d_ff": 32,
+                "patch_size": 4,
+                "individual": True,
+                "use_revin": False,
+            },
+        ]
+
+        for param in params:
+            dm_cfg = {"max_encoder_length": 8, "max_prediction_length": 2}
+            dm_cfg.update(param.get("datamodule_cfg", {}))
+            param["datamodule_cfg"] = dm_cfg
+
+        return params