[ENH] ModernTCN Implemented

Author: Muhammad-Rebaal

pytorch_forecasting/layers/_blocks/_modern_tcn_block.py

@@ -1,29 +1,207 @@
 """
-ModernTCN block: DWConv + ConvFFN + residual.
+ModernTCN building blocks: ReparamLargeKernelConv, ModernTCNBlock, and Flatten_Head.
 """
 
+import torch
 import torch.nn as nn
 
 
 class ModernTCNBlock(nn.Module):
-    def __init__(self, d_model, kernel_size, d_ff, dropout):
+    """
+    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.dwconv = nn.Conv1d(
-            d_model,
-            d_model,
+        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,
-            padding=kernel_size // 2,
-            groups=d_model,
+            stride=1,
+            groups=nvars * d_model,
+            small_kernel_size=small_kernel_size,
         )
         self.norm = nn.BatchNorm1d(d_model)
-        self.pw_conv1 = nn.Conv1d(d_model, d_ff, kernel_size=1)
-        self.pw_conv2 = nn.Conv1d(d_ff, d_model, kernel_size=1)
-        self.act = nn.GELU()
-        self.drop = nn.Dropout(dropout)
+
+        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
+
+
+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)
+
+
+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):
-        residual = x
-        out = self.act(self.norm(self.dwconv(x)))
-        out = self.drop(self.act(self.pw_conv1(out)))
-        out = self.drop(self.pw_conv2(out))
-        return out + residual
+        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/modern_tcn/_modern_tcn_pkg_v2.py

@@ -73,13 +73,21 @@ def get_test_train_params(cls):
                 "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,
+            },
         ]
 
-        default_dm_cfg = {"max_encoder_length": 8, "max_prediction_length": 2}
-
         for param in params:
-            current_dm_cfg = param.get("datamodule_cfg", {})
-            default_dm_cfg.update(current_dm_cfg)
-            param["datamodule_cfg"] = default_dm_cfg
+            dm_cfg = {"max_encoder_length": 8, "max_prediction_length": 2}
+            dm_cfg.update(param.get("datamodule_cfg", {}))
+            param["datamodule_cfg"] = dm_cfg
 
         return params

pytorch_forecasting/models/modern_tcn/_modern_tcn_v2.py

@@ -8,7 +8,10 @@
 from torch.optim import Optimizer
 
 from pytorch_forecasting.layers import RevIN
-from pytorch_forecasting.layers._blocks._modern_tcn_block import ModernTCNBlock
+from pytorch_forecasting.layers._blocks._modern_tcn_block import (
+    Flatten_Head,
+    ModernTCNBlock,
+)
 from pytorch_forecasting.models.base._base_model_v2 import BaseModel
 
 
@@ -27,17 +30,36 @@ class ModernTCN(BaseModel):
     d_model : int
         Embedding dimension per patch.
     kernel_size : int
-        Kernel size for the depthwise convolution.
+        Large kernel size for the reparameterizable depthwise convolution.
+    small_kernel_size : int
+        Small kernel size for the parallel branch in ReparamLargeKernelConv.
     n_blocks : int
-        Number of ModernTCN blocks.
+        Number of ModernTCN encoder blocks.
     d_ff : int
-        Hidden dimension in the pointwise conv FFN.
+        Hidden dimension in the pointwise conv FFNs.
     patch_size : int
-        Number of time steps per patch.
+        Number of time steps per patch (controls stem stride).
     dropout : float
-        Dropout rate.
+        Dropout rate applied inside each encoder block.
+    head_dropout : float
+        Dropout rate applied inside the Flatten_Head projection.
+    individual : bool
+        If True, uses a separate linear projection per variable in the head.
     use_revin : bool
-        Whether to use RevIN normalization.
+        Whether to apply Reversible Instance Normalization before the encoder.
+    logging_metrics : list[nn.Module] or None
+        Additional metrics to log during training.
+    optimizer : Optimizer or str or None
+        Optimizer or name of optimizer (default: ``"adam"``).
+    optimizer_params : dict or None
+        Additional keyword arguments passed to the optimizer constructor.
+    lr_scheduler : str or None
+        Name of a learning rate scheduler recognised by Lightning.
+    lr_scheduler_params : dict or None
+        Additional keyword arguments passed to the scheduler constructor.
+    metadata : dict or None
+        Dataset metadata injected by the package layer (encoder lengths,
+        target dim, etc.).
     """
 
     @classmethod
@@ -54,10 +76,13 @@ def __init__(
         loss: nn.Module,
         d_model: int = 64,
         kernel_size: int = 51,
+        small_kernel_size: int = 5,
         n_blocks: int = 2,
         d_ff: int = 256,
         patch_size: int = 8,
         dropout: float = 0.1,
+        head_dropout: float = 0.1,
+        individual: bool = False,
         use_revin: bool = True,
         logging_metrics: list[nn.Module] | None = None,
         optimizer: Optimizer | str | None = "adam",
@@ -96,18 +121,31 @@ def __init__(
         if self.use_revin:
             self.revin = RevIN(num_features=self.n_channels)
 
-        self.patch_embed = nn.Linear(patch_size, d_model)
+        self.stem = nn.Sequential(
+            nn.Conv1d(1, d_model, kernel_size=patch_size, stride=patch_size),
+            nn.BatchNorm1d(d_model),
+        )
 
         self.blocks = nn.ModuleList(
             [
-                ModernTCNBlock(d_model, kernel_size, d_ff, dropout)
+                ModernTCNBlock(
+                    d_model,
+                    kernel_size,
+                    small_kernel_size,
+                    d_ff,
+                    self.n_channels,
+                    dropout,
+                )
                 for _ in range(n_blocks)
             ]
         )
 
-        self.head = nn.Linear(
-            self.n_patches * d_model,
-            self.prediction_length * self.n_quantiles,
+        self.head = Flatten_Head(
+            individual=individual,
+            n_vars=self.n_channels,
+            nf=d_model * self.n_patches,
+            target_window=self.prediction_length * self.n_quantiles,
+            head_dropout=head_dropout,
         )
 
     def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
@@ -125,21 +163,24 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         B, L, C = input_data.shape
 
         x_enc = input_data.permute(0, 2, 1)
-        x_enc = x_enc.reshape(B, C, self.n_patches, self.patch_size)
-        x_enc = self.patch_embed(x_enc)
-        x_enc = x_enc.permute(0, 1, 3, 2)
-        x_enc = x_enc.reshape(B * C, self.d_model, self.n_patches)
+        x_enc = x_enc.reshape(B * C, 1, L)
+        x_enc = self.stem(x_enc)
 
+        x_enc = x_enc.reshape(B, C, self.d_model, self.n_patches)
         for block in self.blocks:
             x_enc = block(x_enc)
 
-        x_enc = x_enc.reshape(B, C, self.d_model, self.n_patches)
-        x_enc = x_enc.reshape(B, C, -1)
-
         out = self.head(x_enc)
         out = out.reshape(B, C, self.prediction_length, self.n_quantiles)
         out = out.permute(0, 2, 1, 3)
 
+        if self.use_revin:
+            out = out.permute(0, 1, 3, 2)
+            out = out.reshape(B, self.prediction_length * self.n_quantiles, C)
+            out = self.revin(out, mode="denorm")
+            out = out.reshape(B, self.prediction_length, self.n_quantiles, C)
+            out = out.permute(0, 1, 3, 2)
+
         target_indices = list(range(self.n_cont_features, C))
         out = out[:, :, target_indices, :]