Added BatchNorm2d, LayerNorm, InstanceNorm2d nodes.

modules/pytorch_wrapper/ptn_avg_pool2d.py

@@ -5,6 +5,7 @@ from .utils import str_to_dim
 
 class PtnAvgPool2d:
     """
+    Ptn Avg Pool 2d:
     An avgpool layer.  
 
         Args:

modules/pytorch_wrapper/ptn_batch_norm2d.py

@@ -0,0 +1,80 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+from .utils import str_to_dim
+
+class PtnBatchNorm2d:
+    """
+    Ptn Batch Norm 2d:
+    A normalization model to normalize over the batch and spatial axes for each channel.
+
+    Args:
+        num_features (int): Specifies the size of channel axis.
+            - For an image input of shape `[8, 4, 256, 256]`, specify 4.
+
+        affine (bool): If `True`, applies a learnable scaling factor and bias to the normalized elements.
+
+        track_running_stats (bool): If `True`, use exponential moving average (EMA) to store mean and std to be used in eval for normalization. For training, current batch statistics are always used for mean and std irrespective of this flag.
+            
+        momentum (float): The EMA coefficient used if track_running_stats is set to True. Higher values give more weight to the current batch statistics, while lower values retain more past statistics.
+             
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "num_features": ("INT", {"default":1, "min":1, "max":1e6}),
+                "affine": ("BOOLEAN", {"default": True}),
+                "track_running_stats": ("BOOLEAN", {"default": True}),
+                "momentum": ("FLOAT", {"default":0.1, "min":0, "max":1, "step":1e-6}),
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          num_features: int,
+          affine: bool,
+          track_running_stats: bool,
+          momentum: float
+          ) -> Tuple[nn.Module]:
+        """
+        Instantiates a normalization model to normalize over the batch axis and spacial axes.
+
+        Args:
+            num_features (int): Specifies the size of channel axis.
+                - For an image input of shape `[8, 4, 256, 256]`, specify 4.
+
+            affine (bool): If `True`, applies a learnable scaling factor and bias to the normalized elements.
+
+            track_running_stats (bool): If `True`, use exponential moving average (EMA) to store mean and std to be used in eval for normalization. For training, current batch statistics are always used for mean and std irrespective of this flag.
+             
+            momentum (float): The EMA coefficient used if track_running_stats is set to True. Higher values give more weight to the current batch statistics, while lower values retain more past statistics.
+            
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+
+        with torch.inference_mode(False):
+            model = nn.BatchNorm2d(
+                num_features=num_features,
+                affine=affine,
+                track_running_stats=track_running_stats,
+                momentum=momentum)
+                    
+        return (model,)

modules/pytorch_wrapper/ptn_conv2d.py

@@ -6,6 +6,7 @@ from .utils import str_to_dim
 
 class PtnConv2d:
     """
+    Ptn Conv 2d:
     A convolutional model consisting of a single conv2d layer.  
 
         Args:

modules/pytorch_wrapper/ptn_instance_norm2d.py

@@ -0,0 +1,80 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+from .utils import str_to_dim
+
+class PtnInstanceNorm2d:
+    """
+    Ptn Instance Norm 2d:
+    A normalization model to normalize elements over spatial axes within each channel for each sample.
+
+        Args:
+            num_features (int): Specifies the size of channel axis.
+                - For an image input of shape `[8, 4, 256, 256]`, specify 4.
+
+            affine (bool): If `True`, applies a learnable scaling factor and bias to the normalized elements.
+
+            track_running_stats (bool): If `True`, use exponential moving average (EMA) to store mean and std to be used in eval for normalization. For training, current sample data is always used for mean and std irrespective of this flag.
+             
+            momentum (float): The EMA coefficient used if track_running_stats is set to True. Higher value gives more weight to the current statistics (mean & std).
+             
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "num_features": ("INT", {"default":1, "min":1, "max":1e6}),
+                "affine": ("BOOLEAN", {"default": True}),
+                "track_running_stats": ("BOOLEAN", {"default": True}),
+                "momentum": ("FLOAT", {"default":0.1, "min":0, "max":1, "step":1e-6}),
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          num_features: int,
+          affine: bool,
+          track_running_stats: bool,
+          momentum: float
+          ) -> Tuple[nn.Module]:
+        """
+        A normalization model to normalize elements over spatial axes within each channel for each sample.
+
+        Args:
+            num_features (int): Specifies the size of channel axis.
+                - For an image input of shape `[8, 4, 256, 256]`, specify 4.
+
+            affine (bool): If `True`, applies a learnable scaling factor and bias to the normalized elements.
+
+            track_running_stats (bool): If `True`, use exponential moving average (EMA) to store mean and std to be used in eval for normalization. For training, current sample data is always used for mean and std irrespective of this flag.
+             
+            momentum (float): The EMA coefficient used if track_running_stats is set to True. Higher value gives more weight to the current statistics (mean & std).
+            
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+
+        with torch.inference_mode(False):
+            model = nn.InstanceNorm2d(
+                num_features=num_features,
+                affine=affine,
+                track_running_stats=track_running_stats,
+                momentum=momentum)
+                    
+        return (model,)

modules/pytorch_wrapper/ptn_layer_norm.py

@@ -0,0 +1,82 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+from .utils import str_to_dim
+
+class PtnLayerNorm:
+    """
+    A normalization model to normalize elements over specified axes.
+
+        Args:
+            normalized_shape (list): Specifies the shape of elements to normalize.  
+                - For an input of shape `[8, 1024, 768]` (where axes represent `[batch_size, sequence_length, hidden_dim]`):
+                - `normalize_shape=[768]` normalizes within each token (across the hidden dimension).
+                - `normalize_shape=[1024, 768]` normalizes within each sample (across both sequence and token dimensions).  
+                - For an image input of shape `[8, 4, 256, 256]` (where axes represent `[batch_size, channel, height, width]`):
+                - `normalize_shape=[256, 256]` normalizes within each channel (across spatial dimensions).
+                - `normalize_shape=[4, 256, 256]` normalizes within each sample (across channels and spatial dimensions).
+
+            elementwise_affine (bool): If `True`, applies a learnable scaling factor to the normalized elements.
+
+            bias (bool): If `True` and elementwise_affine is also `True`, applies a learnable bias to the normalized elements.
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "normalized_shape": ("STRING", {"default": "[Specify shape here]"}),
+                "elementwise_affine": ("BOOLEAN", {"default": True}),
+                "bias": ("BOOLEAN", {"default": True}),
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          normalized_shape: str,
+          elementwise_affine: bool,
+          bias: bool
+          ) -> Tuple[nn.Module]:
+        """
+        A normalization model to normalize elements over specified axes.
+
+        Args:
+            normalized_shape (list): Specifies the shape of elements to normalize.  
+                - For an input of shape `[8, 1024, 768]` (where axes represent `[batch_size, sequence_length, hidden_dim]`):
+                - `normalize_shape=[768]` normalizes within each token (across the hidden dimension).
+                - `normalize_shape=[1024, 768]` normalizes within each sample (across both sequence and token dimensions).  
+                - For an image input of shape `[8, 4, 256, 256]` (where axes represent `[batch_size, channel, height, width]`):
+                - `normalize_shape=[256, 256]` normalizes within each channel (across spatial dimensions).
+                - `normalize_shape=[4, 256, 256]` normalizes within each sample (across channels and spatial dimensions).
+
+            elementwise_affine (bool): If `True`, applies a learnable scaling factor to the normalized elements.
+
+            bias (bool): If `True` and elementwise_affine is also `True`, applies a learnable bias to the normalized elements.
+            
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+        shape = str_to_dim(normalized_shape)
+
+        with torch.inference_mode(False):
+            model = nn.LayerNorm(
+                normalized_shape=shape,
+                elementwise_affine=elementwise_affine,
+                bias=bias)
+            
+        return (model,)

modules/pytorch_wrapper/ptn_max_pool2d.py

@@ -5,6 +5,7 @@ from .utils import str_to_dim
 
 class PtnMaxPool2d:
     """
+    Ptn Max Pool 2d:
     A maxpool layer.  
 
         Args: