Separate the reusable layers

Author: Muhammad-Rebaal

pytorch_forecasting/layers/_blocks/_modern_tcn_block.py

@@ -1,10 +1,14 @@
 """
-ModernTCN building blocks: ReparamLargeKernelConv, ModernTCNBlock, and Flatten_Head.
+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):
     """
@@ -89,119 +93,3 @@ def forward(self, x):
         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):
-        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/layers/_convolution/__init__.py

@@ -0,0 +1,3 @@
+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/modern_tcn/_modern_tcn_v2.py

@@ -8,10 +8,8 @@
 from torch.optim import Optimizer
 
 from pytorch_forecasting.layers import RevIN
-from pytorch_forecasting.layers._blocks._modern_tcn_block import (
-    Flatten_Head,
-    ModernTCNBlock,
-)
+from pytorch_forecasting.layers._blocks._modern_tcn_block import ModernTCNBlock
+from pytorch_forecasting.layers._head._flatten_head import Flatten_Head
 from pytorch_forecasting.models.base._base_model_v2 import BaseModel