diff --git a/__init__.py b/__init__.py index c7bf831..cfcd7d6 100644 --- a/__init__.py +++ b/__init__.py @@ -146,9 +146,12 @@ from .modules.pytorch_wrapper.ptf_softmax import PtfSoftmax from .modules.pytorch_wrapper.ptf_softplus import PtfSoftplus from .modules.pytorch_wrapper.ptf_tanh import PtfTanh from .modules.pytorch_wrapper.ptn_avg_pool2d import PtnAvgPool2d +from .modules.pytorch_wrapper.ptn_batch_norm2d import PtnBatchNorm2d from .modules.pytorch_wrapper.ptn_chained_model import PtnChainedModel from .modules.pytorch_wrapper.ptn_conv2d import PtnConv2d from .modules.pytorch_wrapper.ptn_conv_model import PtnConvModel +from .modules.pytorch_wrapper.ptn_instance_norm2d import PtnInstanceNorm2d +from .modules.pytorch_wrapper.ptn_layer_norm import PtnLayerNorm from .modules.pytorch_wrapper.ptn_linear import PtnLinear from .modules.pytorch_wrapper.ptn_linear_model import PtnLinearModel from .modules.pytorch_wrapper.ptn_max_pool2d import PtnMaxPool2d @@ -324,9 +327,12 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtfSoftplus": PtfSoftplus, "PtfTanh": PtfTanh, "PtnAvgPool2d": PtnAvgPool2d, + "PtnBatchNorm2d": PtnBatchNorm2d, "PtnChainedModel": PtnChainedModel, "PtnConv2d": PtnConv2d, "PtnConvModel": PtnConvModel, + "PtnInstanceNorm2d": PtnInstanceNorm2d, + "PtnLayerNorm": PtnLayerNorm, "PtnLinear": PtnLinear, "PtnLinearModel": PtnLinearModel, "PtnMaxPool2d": PtnMaxPool2d, @@ -501,13 +507,16 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtfSoftmax": "Ptf Softmax", "PtfSoftplus": "Ptf Softplus", "PtfTanh": "Ptf Tanh", - "PtnAvgPool2d": "Ptn Avg Pool2d", + "PtnAvgPool2d": "Ptn Avg Pool 2d", + "PtnBatchNorm2d": "Ptn Batch Norm 2d", "PtnChainedModel": "Ptn Chained Model", - "PtnConv2d": "Ptn Conv2d", + "PtnConv2d": "Ptn Conv 2d", "PtnConvModel": "Ptn Conv Model", + "PtnInstanceNorm2d": "Ptn Instance Norm 2d", + "PtnLayerNorm": "Ptn Layer Norm", "PtnLinear": "Ptn Linear", "PtnLinearModel": "Ptn Linear Model", - "PtnMaxPool2d": "Ptn Max Pool2d", + "PtnMaxPool2d": "Ptn Max Pool 2d", "PtnPreAddChannelAxis": "Ptn Pre Add Channel Axis", "PtnPreFlatten": "Ptn Pre Flatten", "PtnResnetModel": "Ptn Resnet Model", diff --git a/dev/scripts/init_py_generate.py b/dev/scripts/init_py_generate.py index 32c511b..07d1f9b 100644 --- a/dev/scripts/init_py_generate.py +++ b/dev/scripts/init_py_generate.py @@ -90,7 +90,8 @@ def format_display_name(class_name): ("MPL Scatter", "MPL Scatter Plot"), ("SNS Bar", "SNS Bar Chart"), ("SNS Line", "SNS Line Plot"), - ("SNS Scatter", "SNS Scatter Plot") + ("SNS Scatter", "SNS Scatter Plot"), + ("2d", " 2d"), ] for p in replace_pairs: final_result = final_result.replace(p[0], p[1]) diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index b572151..076e81f 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -138,13 +138,16 @@ ## PyTorch wrapper - Model | Node | Description | | --- | --- | -| [Ptn Avg Pool](ptn_avg_pool2d.md) | An avgpool layer. | +| [Ptn Avg Pool 2d](ptn_avg_pool2d.md) | An avgpool layer. | +| [Ptn Batch Norm 2d](ptn_batch_norm2d.md) | A normalization model to normalize over the batch and spatial axes for each channel. | | [Ptn Chained Model](ptn_chained_model.md) | Constructs a chained PyTorch model. | -| [Ptn Conv](ptn_conv2d.md) | A convolutional model consisting of a single conv2d layer. | +| [Ptn Conv 2d](ptn_conv2d.md) | A convolutional model consisting of a single conv2d layer. | | [Ptn Conv Model](ptn_conv_model.md) | A convolutional model consisting of multiple convolutional layers. | +| [Ptn Instance Norm 2d](ptn_instance_norm2d.md) | A normalization model to normalize elements over spatial axes within each channel for each sample. | +| [Ptn Layer Norm](ptn_layer_norm.md) | A normalization model to normalize elements over specified axes. | | [Ptn Linear](ptn_linear.md) | A linear model consisting of a single dense layer. | | [Ptn Linear Model](ptn_linear_model.md) | A linear model consisting of dense layers. | -| [Ptn Max Pool](ptn_max_pool2d.md) | A maxpool layer. | +| [Ptn Max Pool 2d](ptn_max_pool2d.md) | A maxpool layer. | | [Ptn Pre Add Channel Axis](ptn_pre_add_channel_axis.md) | Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w) | | [Ptn Pre Flatten](ptn_pre_flatten.md) | Flattens the input tensor before processing the tensor in the specified model. | | [Ptn Resnet Model](ptn_resnet_model.md) | A Resnet model consisting of multiple Resnet layers. | diff --git a/docs/reference/ptn_avg_pool2d.md b/docs/reference/ptn_avg_pool2d.md index c13846b..7999888 100644 --- a/docs/reference/ptn_avg_pool2d.md +++ b/docs/reference/ptn_avg_pool2d.md @@ -1,4 +1,4 @@ -# Ptn Avg Pool +# Ptn Avg Pool 2d An avgpool layer. Args: diff --git a/docs/reference/ptn_batch_norm2d.md b/docs/reference/ptn_batch_norm2d.md new file mode 100644 index 0000000..82df011 --- /dev/null +++ b/docs/reference/ptn_batch_norm2d.md @@ -0,0 +1,30 @@ +# 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. + +## Input +| Name | Data type | +|---|---| +| num_features | Int | +| affine | Boolean | +| track_running_stats | Boolean | +| momentum | Float | + +## Output +| Data type | +|---| +| Ptmodel | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_conv2d.md b/docs/reference/ptn_conv2d.md index 55d574c..6298457 100644 --- a/docs/reference/ptn_conv2d.md +++ b/docs/reference/ptn_conv2d.md @@ -1,4 +1,4 @@ -# Ptn Conv +# Ptn Conv 2d A convolutional model consisting of a single conv2d layer. Args: diff --git a/docs/reference/ptn_instance_norm2d.md b/docs/reference/ptn_instance_norm2d.md new file mode 100644 index 0000000..c1077b3 --- /dev/null +++ b/docs/reference/ptn_instance_norm2d.md @@ -0,0 +1,30 @@ +# 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). + +## Input +| Name | Data type | +|---|---| +| num_features | Int | +| affine | Boolean | +| track_running_stats | Boolean | +| momentum | Float | + +## Output +| Data type | +|---| +| Ptmodel | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_layer_norm.md b/docs/reference/ptn_layer_norm.md new file mode 100644 index 0000000..a73ef73 --- /dev/null +++ b/docs/reference/ptn_layer_norm.md @@ -0,0 +1,32 @@ +# Ptn Layer Norm +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. + +## Input +| Name | Data type | +|---|---| +| normalized_shape | String | +| elementwise_affine | Boolean | +| bias | Boolean | + +## Output +| Data type | +|---| +| Ptmodel | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_max_pool2d.md b/docs/reference/ptn_max_pool2d.md index 544a1cc..fc2ba34 100644 --- a/docs/reference/ptn_max_pool2d.md +++ b/docs/reference/ptn_max_pool2d.md @@ -1,4 +1,4 @@ -# Ptn Max Pool +# Ptn Max Pool 2d A maxpool layer. Args: diff --git a/modules/pytorch_wrapper/ptn_avg_pool2d.py b/modules/pytorch_wrapper/ptn_avg_pool2d.py index 0197076..a66a680 100644 --- a/modules/pytorch_wrapper/ptn_avg_pool2d.py +++ b/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: diff --git a/modules/pytorch_wrapper/ptn_batch_norm2d.py b/modules/pytorch_wrapper/ptn_batch_norm2d.py new file mode 100644 index 0000000..62f501b --- /dev/null +++ b/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,) diff --git a/modules/pytorch_wrapper/ptn_conv2d.py b/modules/pytorch_wrapper/ptn_conv2d.py index ed780f0..648968c 100644 --- a/modules/pytorch_wrapper/ptn_conv2d.py +++ b/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: diff --git a/modules/pytorch_wrapper/ptn_instance_norm2d.py b/modules/pytorch_wrapper/ptn_instance_norm2d.py new file mode 100644 index 0000000..7a5cdee --- /dev/null +++ b/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,) diff --git a/modules/pytorch_wrapper/ptn_layer_norm.py b/modules/pytorch_wrapper/ptn_layer_norm.py new file mode 100644 index 0000000..24f0275 --- /dev/null +++ b/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,) diff --git a/modules/pytorch_wrapper/ptn_max_pool2d.py b/modules/pytorch_wrapper/ptn_max_pool2d.py index ddc3866..760bc88 100644 --- a/modules/pytorch_wrapper/ptn_max_pool2d.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index 6c551c3..607d206 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "pt-wrapper" description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training." -version = "3.1.3" +version = "3.1.4" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/ptn_batch_norm2d_test.py b/tests/pytorch_wrapper/ptn_batch_norm2d_test.py new file mode 100644 index 0000000..5618cef --- /dev/null +++ b/tests/pytorch_wrapper/ptn_batch_norm2d_test.py @@ -0,0 +1,64 @@ +import os +import sys +import unittest +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) +from pytorch_wrapper.ptn_batch_norm2d import PtnBatchNorm2d + + +class TestPtnBatchNorm2d(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnBatchNorm2d() + + def check_normalization(self, x, num_features): + model = self.model_node.f(num_features, + affine=False, + track_running_stats=False, + momentum=0.001)[0] + + out = model(x) + self.assertEqual(out.shape, x.shape) + + # Compute std along normalized axes before and after + std_before = x.std(dim=(0, 2, 3), unbiased=True) + std_after = out.std(dim=(0, 2, 3), unbiased=True) + mean_after = out.mean(dim=(0, 2, 3)) + + # Ensure std is not one before normalization + self.assertFalse( + torch.allclose( + std_before, + torch.ones_like(std_after), + atol=1e-1)) + + # Ensure normalization happened (std should be 1 unless elements are very close) + self.assertTrue( + torch.allclose( + std_after, + torch.ones_like(std_after), + atol=1e-1)) + + self.assertTrue( + torch.allclose( + mean_after, + torch.zeros_like(mean_after), + atol=1e-1)) + + def test_normalization_cases(self): + test_cases = [ + # Data, num_features on the channel axis + (torch.randn(2, 3, 3, 3) * 2, 3), # c = 3 + (torch.randn(3, 3, 3, 3) * 2, 3), # c = 3 + ] + + for x, norm_shape in test_cases: + with self.subTest(num_features=norm_shape): + self.check_normalization(x, norm_shape) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptn_instance_norm2d_test.py b/tests/pytorch_wrapper/ptn_instance_norm2d_test.py new file mode 100644 index 0000000..f6a2e0e --- /dev/null +++ b/tests/pytorch_wrapper/ptn_instance_norm2d_test.py @@ -0,0 +1,63 @@ +import os +import sys +import unittest +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) +from pytorch_wrapper.ptn_instance_norm2d import PtnInstanceNorm2d + + +class TestPtnInstanceNorm2d(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnInstanceNorm2d() + + def check_normalization(self, x, num_features): + model = self.model_node.f(num_features, True, True, 0.001)[0] + out = model(x) + self.assertEqual(out.shape, x.shape) + + # Compute std along normalized axes before and after + std_before = x.std(dim=(-2,-1), unbiased=False) + std_after = out.std(dim=(-2,-1), unbiased=False) + + mean_after = out.mean(dim=(-2,-1)) + + # Ensure std is not one before normalization + self.assertFalse( + torch.allclose( + std_before, + torch.ones_like(std_after), + atol=1e-3)) + + # Ensure normalization happened (std should be 1 unless elements are very close) + self.assertTrue( + torch.allclose( + std_after, + torch.ones_like(std_after), + atol=1e-3)) + + self.assertTrue( + torch.allclose( + mean_after, + torch.zeros_like(mean_after), + atol=1e-3)) + + def test_normalization_cases(self): + test_cases = [ + # Data, num_features on the channel axis + (torch.randn(3, 2, 2) * 2, 3), # c = 3 + (torch.randn(1, 3, 2, 2) * 2, 3), # c = 3 + (torch.randn(2, 3, 2, 2) * 2, 3), # c = 3 + (torch.randn(8, 3, 2, 2) * 2, 3), # c = 3 + ] + + for x, norm_shape in test_cases: + with self.subTest(num_features=norm_shape): + self.check_normalization(x, norm_shape) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptn_layer_norm_test.py b/tests/pytorch_wrapper/ptn_layer_norm_test.py new file mode 100644 index 0000000..983f3a3 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_layer_norm_test.py @@ -0,0 +1,70 @@ +import os +import sys +import unittest +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) +from pytorch_wrapper.ptn_layer_norm import PtnLayerNorm + + +class TestPtnLayerNorm(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnLayerNorm() + + def check_normalization(self, x, normalized_shape): + model = self.model_node.f(str(normalized_shape), True, True)[0] + out = model(x) + self.assertEqual(out.shape, x.shape) + + # Compute std along normalized axes before and after + std_before = x.std(dim=tuple( + range(-len(normalized_shape), 0) + ), unbiased=False) + std_after = out.std(dim=tuple( + range(-len(normalized_shape), 0) # (-2, -1) for last two or (-3, -2, -1) for last three + ), unbiased=False) + + mean_after = out.mean(dim=tuple( + range(-len(normalized_shape), 0) # (-2, -1) for last two or (-3, -2, -1) for last three + )) + + # Ensure std is not one before normalization + self.assertFalse( + torch.allclose( + std_before, + torch.ones_like(std_after), + atol=1e-3)) + + # Ensure normalization happened (std should be 1 unless elements are very close) + self.assertTrue( + torch.allclose( + std_after, + torch.ones_like(std_after), + atol=1e-3)) + + self.assertTrue( + torch.allclose( + mean_after, + torch.zeros_like(mean_after), + atol=1e-3)) + + def test_normalization_cases(self): + test_cases = [ + # Data, normalized axis + (torch.randn(1, 2, 4) * 2, [4]), # token + (torch.randn(1, 2, 4) * 2, [2, 4]), # seq and token => per sample + (torch.randn(2, 2, 4) * 2, [2, 4]), # seq and token => per sample + (torch.randn(1, 3, 2, 2) * 2, [2, 2]), # h, w + (torch.randn(1, 3, 2, 2) * 2, [3, 2, 2]), # c, h, w => per sample + ] + + for x, norm_shape in test_cases: + with self.subTest(normalized_shape=norm_shape): + self.check_normalization(x, norm_shape) + + +if __name__ == "__main__": + unittest.main()