diff --git a/.gitignore b/.gitignore index e9242fa..0c21420 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,5 @@ __pycache__/ models/ *.pth /datasets/ -examples/datasets/.~lock* \ No newline at end of file +examples/datasets/.~lock* +.vscode \ No newline at end of file diff --git a/__init__.py b/__init__.py index fb852d5..e9091a0 100644 --- a/__init__.py +++ b/__init__.py @@ -121,10 +121,13 @@ from .modules.pytorch_wrapper.ptf_silu import PtfSiLU 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_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_linear import PtnLinear from .modules.pytorch_wrapper.ptn_linear_model import PtnLinearModel +from .modules.pytorch_wrapper.ptn_max_pool2d import PtnMaxPool2d from .modules.pytorch_wrapper.ptn_resnet_model import PtnResnetModel from .modules.pytorch_wrapper.pto_adam import PtoAdam from .modules.pytorch_wrapper.pto_adamw import PtoAdamW @@ -270,10 +273,13 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtfSoftmax": PtfSoftmax, "PtfSoftplus": PtfSoftplus, "PtfTanh": PtfTanh, + "PtnAvgPool2d": PtnAvgPool2d, "PtnChainedModel": PtnChainedModel, + "PtnConv2d": PtnConv2d, "PtnConvModel": PtnConvModel, "PtnLinear": PtnLinear, "PtnLinearModel": PtnLinearModel, + "PtnMaxPool2d": PtnMaxPool2d, "PtnResnetModel": PtnResnetModel, "PtoAdam": PtoAdam, "PtoAdamW": PtoAdamW, @@ -419,10 +425,13 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtfSoftmax": "Ptf Softmax", "PtfSoftplus": "Ptf Softplus", "PtfTanh": "Ptf Tanh", + "PtnAvgPool2d": "Ptn Avg Pool2d", "PtnChainedModel": "Ptn Chained Model", + "PtnConv2d": "Ptn Conv2d", "PtnConvModel": "Ptn Conv Model", "PtnLinear": "Ptn Linear", "PtnLinearModel": "Ptn Linear Model", + "PtnMaxPool2d": "Ptn Max Pool2d", "PtnResnetModel": "Ptn Resnet Model", "PtoAdam": "Pto Adam", "PtoAdamW": "Pto AdamW", diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index 232a4a0..7ffd658 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -110,6 +110,17 @@ | [Pt Einsum](pt_einsum.md) | Performs Tensor operations specified in the Einstein summation equation. | | [Pt Mat Mul](pt_matmul.md) | Performs matrix multiplication on two PyTorch tensors. | | [Pt Mm](pt_mm.md) | Performs 2D matrix multiplication on two PyTorch tensors. | +## PyTorch wrapper - Model +| Node | Description | +| --- | --- | +| [Ptn Avg Pool](ptn_avg_pool2d.md) | An avgpool layer. | +| [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 Model](ptn_conv_model.md) | A convolutional model consisting of multiple convolutional layers. | +| [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 Resnet Model](ptn_resnet_model.md) | A Resnet model consisting of multiple Resnet layers. | ## PyTorch wrapper - Optimizer | Node | Description | | --- | --- | @@ -172,11 +183,6 @@ | [Pt Train Classification Model](pt_train_classification_model.md) | Trains a classification model using a given dataset, optimizer, and number of epochs. | | [Pt Train Classification Model Lr](pt_train_classification_model_lr.md) | Trains a classification model using a given dataset, optimizer, and number of epochs with learning rate decay. | | [Pt Train Regression Model](pt_train_regression_model.md) | Trains a regression model using a given dataset, optimizer, and number of epochs. | -| [Ptn Chained Model](ptn_chained_model.md) | Constructs a chained PyTorch model. | -| [Ptn Conv Model](ptn_conv_model.md) | A convolutional model consisting of multiple convolutional layers. | -| [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 Resnet Model](ptn_resnet_model.md) | A Resnet model consisting of multiple Resnet layers. | | [Pto Lr Scheduler Cosine Annealing](pto_lr_scheduler_cosine_annealing.md) | Creates a cosine annealing learning rate scheduler for an optimizer. | | [Pto Lr Scheduler Reduce On Plateau](pto_lr_scheduler_reduce_on_plateau.md) | Creates a reduce-on-plateau learning rate scheduler for an optimizer. | | [Pto Lr Scheduler Step](pto_lr_scheduler_step.md) | Creates a StepLR learning rate scheduler for an optimizer. | diff --git a/docs/reference/ptn_avg_pool2d.md b/docs/reference/ptn_avg_pool2d.md new file mode 100644 index 0000000..c13846b --- /dev/null +++ b/docs/reference/ptn_avg_pool2d.md @@ -0,0 +1,24 @@ +# Ptn Avg Pool +An avgpool layer. + + Args: + kernel_size (str): The size of the kernel + stride (str): The stride to be used for sliding the kernel + padding (str): The amount of padding added to the input + +## Input +| Name | Data type | +|---|---| +| kernel_size | String | +| stride | String | +| padding | String | + +## 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_chained_model.md b/docs/reference/ptn_chained_model.md index ca2ae03..039c6fb 100644 --- a/docs/reference/ptn_chained_model.md +++ b/docs/reference/ptn_chained_model.md @@ -19,6 +19,6 @@ Constructs a chained PyTorch model. | Ptmodel |
-Category: PyTorch wrapper - Training +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 new file mode 100644 index 0000000..e5a0cd4 --- /dev/null +++ b/docs/reference/ptn_conv2d.md @@ -0,0 +1,38 @@ +# Ptn Conv +A convolutional model consisting of a single conv2d layer. + + Args: + in_channels (int): The number of input channels + out_channels (int): The number of output channels + kernel_size (str): The size of the kernel + stride (str): The stride to be used for sliding the kernel + padding (str): The amount of padding added to the input + dilation (str): The distance between adjacent elements in the kernel to adjust receptive fields. + groups (str): The number of groups used to divide input and output channels for separate processing. + bias (bool): If `True`, adds a learnable bias to the output + padding_mode (str): Specifies how padding values are set in the input before convolution. +t. + bias (bool): Use bias or not. + +## Input +| Name | Data type | +|---|---| +| in_channels | Int | +| out_channels | Int | +| kernel_size | String | +| stride | String | +| padding | String | +| dilation | String | +| groups | Int | +| bias | Boolean | +| padding_mode | Zeros | + +## 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_conv_model.md b/docs/reference/ptn_conv_model.md index 44687ae..80c9d05 100644 --- a/docs/reference/ptn_conv_model.md +++ b/docs/reference/ptn_conv_model.md @@ -40,6 +40,6 @@ A convolutional model consisting of multiple convolutional layers. | Ptmodel |
-Category: PyTorch wrapper - Training +Category: PyTorch wrapper - Model ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_linear.md b/docs/reference/ptn_linear.md index 2f43cd2..967addd 100644 --- a/docs/reference/ptn_linear.md +++ b/docs/reference/ptn_linear.md @@ -19,6 +19,6 @@ A linear model consisting of a single dense layer. | Ptmodel |
-Category: PyTorch wrapper - Training +Category: PyTorch wrapper - Model ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_linear_model.md b/docs/reference/ptn_linear_model.md index f043c93..cced261 100644 --- a/docs/reference/ptn_linear_model.md +++ b/docs/reference/ptn_linear_model.md @@ -30,6 +30,6 @@ A linear model consisting of dense layers. | Ptmodel |
-Category: PyTorch wrapper - Training +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 new file mode 100644 index 0000000..544a1cc --- /dev/null +++ b/docs/reference/ptn_max_pool2d.md @@ -0,0 +1,26 @@ +# Ptn Max Pool +A maxpool layer. + + Args: + kernel_size (str): The size of the kernel + stride (str): The stride to be used for sliding the kernel + padding (str): The amount of padding added to the input + dilation (str): The distance between adjacent elements in the kernel to adjust receptive fields. + +## Input +| Name | Data type | +|---|---| +| kernel_size | String | +| stride | String | +| padding | String | +| dilation | String | + +## 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_resnet_model.md b/docs/reference/ptn_resnet_model.md index 6ac65c9..f5ff171 100644 --- a/docs/reference/ptn_resnet_model.md +++ b/docs/reference/ptn_resnet_model.md @@ -23,6 +23,6 @@ A Resnet model consisting of multiple Resnet layers. | Ptmodel |
-Category: PyTorch wrapper - Training +Category: PyTorch wrapper - Model ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/modules/pytorch_wrapper/ptn_avg_pool2d.py b/modules/pytorch_wrapper/ptn_avg_pool2d.py new file mode 100644 index 0000000..0197076 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_avg_pool2d.py @@ -0,0 +1,68 @@ +from typing import Any, Dict, Tuple +import torch +import torch.nn as nn +from .utils import str_to_dim + +class PtnAvgPool2d: + """ + An avgpool layer. + + Args: + kernel_size (str): The size of the kernel + stride (str): The stride to be used for sliding the kernel + padding (str): The amount of padding added to the input + + 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": { + "kernel_size": ("STRING", {"default":"2"}), + "stride": ("STRING", {"default":"2"}), + "padding": ("STRING", {"default":"0"}) + } + } + + RETURN_TYPES: Tuple[str] = ("PTMODEL",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, + kernel_size: str, + stride: str, + padding: str + ) -> Tuple[nn.Module]: + """ + Constructs a dense model. + + Args: + kernel_size (str): The size of the kernel + stride (str): The stride to be used for sliding the kernel + padding (str): The amount of padding added to the input + + Returns: + Tuple[nn.Module]: A tuple containing the instantiated PyTorch model. + """ + k = str_to_dim(kernel_size) + s = str_to_dim(stride) + p = str_to_dim(padding) + + with torch.inference_mode(False): + model = nn.AvgPool2d( + kernel_size=k, + stride=s, + padding=p) + + return (model,) diff --git a/modules/pytorch_wrapper/ptn_chained_model.py b/modules/pytorch_wrapper/ptn_chained_model.py index 8f7adbf..8e27175 100644 --- a/modules/pytorch_wrapper/ptn_chained_model.py +++ b/modules/pytorch_wrapper/ptn_chained_model.py @@ -59,7 +59,7 @@ class PtnChainedModel: - model_b: Second model in the chain. - closure: Optional differentiable function to be called at the end of forward. - category: PyTorch wrapper - Training + category: PyTorch wrapper - Model """ @classmethod diff --git a/modules/pytorch_wrapper/ptn_conv2d.py b/modules/pytorch_wrapper/ptn_conv2d.py new file mode 100644 index 0000000..ed780f0 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_conv2d.py @@ -0,0 +1,109 @@ +from typing import Any, Dict, Tuple +import torch +import torch.nn as nn +import ast +from .utils import str_to_dim + +class PtnConv2d: + """ + A convolutional model consisting of a single conv2d layer. + + Args: + in_channels (int): The number of input channels + out_channels (int): The number of output channels + kernel_size (str): The size of the kernel + stride (str): The stride to be used for sliding the kernel + padding (str): The amount of padding added to the input + dilation (str): The distance between adjacent elements in the kernel to adjust receptive fields. + groups (str): The number of groups used to divide input and output channels for separate processing. + bias (bool): If `True`, adds a learnable bias to the output + padding_mode (str): Specifies how padding values are set in the input before convolution. +t. + bias (bool): Use bias or not. + + 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": { + "in_channels": ("INT", {"default": 1, "min": 1, "max": 1e6}), + "out_channels": ("INT", {"default": 1, "min": 1, "max": 1e6}), + "kernel_size": ("STRING", {"default":"3", "multiline": False}), + "stride": ("STRING", {"default":"1", "multiline": False}), + "padding": ("STRING", {"default":"same", "multiline": False}), + "dilation": ("STRING", {"default":"1", "multiline": False}), + "groups": ("INT", {"default": 1, "min": 1, "max": 1e6}), + "bias": ("BOOLEAN", {"default": True}), + "padding_mode": (("zeros", "reflect", "replicate", "circular"),) + } + } + + RETURN_TYPES: Tuple[str] = ("PTMODEL",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, + in_channels: int, + out_channels: int, + kernel_size: str, + stride: str, + padding: str, + dilation: str, + groups: str, + bias: bool, + padding_mode: str + ) -> Tuple[nn.Module]: + """ + Constructs a dense model. + + Args: + in_channels (int): The number of input channels + out_channels (int): The number of output channels + kernel_size (str): The size of the kernel + stride (str): The stride to be used for sliding the kernel + padding (str): The amount of padding added to the input + dilation (str): The distance between adjacent elements in the kernel to adjust receptive fields. + groups (str): The number of groups used to divide input and output channels for separate processing. + bias (bool): If `True`, adds a learnable bias to the output + padding_mode (str): Specifies how padding values are set in the input before convolution. +t. + bias (bool): Use bias or not. + + Returns: + Tuple[nn.Module]: A tuple containing the instantiated PyTorch model. + """ + k = str_to_dim(kernel_size) + s = str_to_dim(stride) + if(padding == "same"): + p = padding + elif(padding=="valid"): + p = padding + else: + p = str_to_dim(padding) + d = str_to_dim(dilation) + + with torch.inference_mode(False): + model = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=k, + stride=s, + padding=p, + dilation=d, + groups=groups, + bias=bias, + padding_mode=padding_mode) + + return (model,) diff --git a/modules/pytorch_wrapper/ptn_conv_model.py b/modules/pytorch_wrapper/ptn_conv_model.py index a31fe36..0eac08d 100644 --- a/modules/pytorch_wrapper/ptn_conv_model.py +++ b/modules/pytorch_wrapper/ptn_conv_model.py @@ -30,7 +30,7 @@ class PtnConvModel: - `downsample_list`: A string representing a Python list of boolean values indicating whether to downsample at each layer. Example: `"[True,True,True,True,False]"`. - category: PyTorch wrapper - Training + category: PyTorch wrapper - Model """ @classmethod diff --git a/modules/pytorch_wrapper/ptn_linear.py b/modules/pytorch_wrapper/ptn_linear.py index 19671b4..c85e89a 100644 --- a/modules/pytorch_wrapper/ptn_linear.py +++ b/modules/pytorch_wrapper/ptn_linear.py @@ -12,7 +12,7 @@ class PtnLinear: - `out_features`: The number of output features. - `bias`: Use bias or not. - category: PyTorch wrapper - Training + category: PyTorch wrapper - Model """ @classmethod diff --git a/modules/pytorch_wrapper/ptn_linear_model.py b/modules/pytorch_wrapper/ptn_linear_model.py index bd2108e..d201ffb 100644 --- a/modules/pytorch_wrapper/ptn_linear_model.py +++ b/modules/pytorch_wrapper/ptn_linear_model.py @@ -25,7 +25,7 @@ class PtnLinearModel: - `num_layers`: The number of layers in the model. - category: PyTorch wrapper - Training + category: PyTorch wrapper - Model """ @classmethod @@ -38,8 +38,8 @@ class PtnLinearModel: """ return { "required": { - "dim_list": ("STRING", {"default": "784,10"}), - "bias_list": ("STRING", {"default": "{True}"}), + "dim_list": ("STRING", {"default": "[784,10]"}), + "bias_list": ("STRING", {"default": "[True]"}), "num_layers": ("INT", {"default": 1, "min": 1, "max": 2000}), } } diff --git a/modules/pytorch_wrapper/ptn_max_pool2d.py b/modules/pytorch_wrapper/ptn_max_pool2d.py new file mode 100644 index 0000000..ddc3866 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_max_pool2d.py @@ -0,0 +1,75 @@ +from typing import Any, Dict, Tuple +import torch +import torch.nn as nn +from .utils import str_to_dim + +class PtnMaxPool2d: + """ + A maxpool layer. + + Args: + kernel_size (str): The size of the kernel + stride (str): The stride to be used for sliding the kernel + padding (str): The amount of padding added to the input + dilation (str): The distance between adjacent elements in the kernel to adjust receptive fields. + + + 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": { + "kernel_size": ("STRING", {"default":"2"}), + "stride": ("STRING", {"default":"2"}), + "padding": ("STRING", {"default":"0"}), + "dilation": ("STRING", {"default":"1"}) + } + } + + RETURN_TYPES: Tuple[str] = ("PTMODEL",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, + kernel_size: str, + stride: str, + padding: str, + dilation: str + ) -> Tuple[nn.Module]: + """ + Constructs a dense model. + + Args: + kernel_size (str): The size of the kernel + stride (str): The stride to be used for sliding the kernel + padding (str): The amount of padding added to the input + dilation (str): The distance between adjacent elements in the kernel to adjust receptive fields. + + Returns: + Tuple[nn.Module]: A tuple containing the instantiated PyTorch model. + """ + k = str_to_dim(kernel_size) + s = str_to_dim(stride) + p = str_to_dim(padding) + d = str_to_dim(dilation) + + with torch.inference_mode(False): + model = nn.MaxPool2d( + kernel_size=k, + stride=s, + padding=p, + dilation=d) + + return (model,) diff --git a/modules/pytorch_wrapper/ptn_resnet_model.py b/modules/pytorch_wrapper/ptn_resnet_model.py index 9638fcb..97aa777 100644 --- a/modules/pytorch_wrapper/ptn_resnet_model.py +++ b/modules/pytorch_wrapper/ptn_resnet_model.py @@ -17,7 +17,7 @@ class PtnResnetModel: - `num_blocks`: Number of Resnet blocks for 64 channel blocks. Same number of blocks will be created for 128 channel and 256 channel blocks. - category: PyTorch wrapper - Training + category: PyTorch wrapper - Model """ @classmethod diff --git a/pyproject.toml b/pyproject.toml index 3b30382..2802d5b 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.0.11" +version = "3.0.12" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/ptn_avg_pool2d_test.py b/tests/pytorch_wrapper/ptn_avg_pool2d_test.py new file mode 100644 index 0000000..c399492 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_avg_pool2d_test.py @@ -0,0 +1,52 @@ +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_avg_pool2d import PtnAvgPool2d + + +class TestPtnAvgPool2dModel(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnAvgPool2d() + + def test_kernel_stride_padding_as_int(self): + self.model = self.model_node.f( + "2", # kernel_size (string) + "2", # stride (string) + "0" # padding (string) + )[0] + x = torch.ones(1, 3, 64, 64) + out = self.model(x) + self.assertEqual(out.size(), torch.Size([1, 3, 32, 32])) + + def test_kernel_stride_padding_as_tuple(self): + self.model = self.model_node.f( + "(2, 2)", # kernel_size (tuple as string) + "(2, 2)", # stride (tuple as string) + "(1, 1)" # padding (tuple as string) + )[0] + x = torch.ones(1, 3, 64, 64) + out = self.model(x) + self.assertEqual(out.size(), torch.Size([1, 3, 33, 33])) + + def test_padding_as_different_values(self): + for padding in ["0", "1"]: + with self.subTest(padding=padding): + self.model = self.model_node.f( + "2", # kernel_size (string) + "2", # stride (string) + padding # padding (string) + )[0] + x = torch.ones(1, 3, 64, 64) + out = self.model(x) + expected_size = 32 + int(padding) # Adjust for padding effect + self.assertEqual(out.size(), torch.Size([1, 3, expected_size, expected_size])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptn_conv2d_test.py b/tests/pytorch_wrapper/ptn_conv2d_test.py new file mode 100644 index 0000000..fb030d3 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_conv2d_test.py @@ -0,0 +1,104 @@ +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_conv2d import PtnConv2d + + +class TestPtnConv2d(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnConv2d() + + def test_kernel_stride_padding_as_int(self): + self.model = self.model_node.f( + 3, # in_channels + 32, # out_channels + "3", # kernel_size (int) + "1", # stride (int) + "1", # padding (int) + "1", # dilation + 1, # groups + True, # bias + "zeros" # padding_mode + )[0] + x = torch.ones(1, 3, 64, 64) + out = self.model(x) + self.assertEqual(out.size(), torch.Size([1, 32, 64, 64])) + param_names = list(map(lambda e: e[0], self.model.named_parameters())) + self.assertTrue("bias" in param_names) + + def test_kernel_stride_padding_as_tuple(self): + self.model = self.model_node.f( + 3, # in_channels + 1, # out_channels + "(3, 3)", # kernel_size (tuple) + "(1, 1)", # stride (tuple) + "(1, 1)", # padding (tuple) + "(1, 1)", # dilation (tuple) + 1, # groups + False, # bias + "zeros" # padding_mode + )[0] + x = torch.ones(1, 3, 64, 64) + out = self.model(x) + self.assertEqual(out.size(), torch.Size([1, 1, 64, 64])) + param_names = list(map(lambda e: e[0], self.model.named_parameters())) + self.assertFalse("bias" in param_names) + + def test_padding_as_same_valid(self): + for padding_mode in ["zeros", "reflect", "replicate", "circular"]: + with self.subTest(padding_mode=padding_mode): + self.model = self.model_node.f( + 3, # in_channels + 16, # out_channels + "3", # kernel_size (string) + "1", # stride (string) + "same", # padding (string "same") + "1", # dilation (string) + 1, # groups + True, # bias + padding_mode # padding_mode + )[0] + x = torch.ones(1, 3, 64, 64) + out = self.model(x) + self.assertEqual(out.size(), torch.Size([1, 16, 64, 64])) + + self.model = self.model_node.f( + 3, # in_channels + 16, # out_channels + "3", # kernel_size (string) + "1", # stride (string) + "valid", # padding (string "valid") + "1", # dilation (string) + 1, # groups + True, # bias + padding_mode # padding_mode + )[0] + x = torch.ones(1, 3, 64, 64) + out = self.model(x) + self.assertEqual(out.size(), torch.Size([1, 16, 62, 62])) + + def test_groups_2(self): + self.model = self.model_node.f( + 4, # in_channels (divisible by groups) + 16, # out_channels + "3", # kernel_size + "1", # stride + "same", # padding + "1", # dilation + 2, # groups + True, # bias + "zeros" # padding_mode + )[0] + x = torch.ones(1, 4, 64, 64) + out = self.model(x) + self.assertEqual(out.size(), torch.Size([1, 16, 64, 64])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptn_max_pool2d_test.py b/tests/pytorch_wrapper/ptn_max_pool2d_test.py new file mode 100644 index 0000000..5ff4274 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_max_pool2d_test.py @@ -0,0 +1,55 @@ +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_max_pool2d import PtnMaxPool2d + + +class TestPtnMaxPool2dModel(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnMaxPool2d() + + def test_kernel_stride_padding_as_int(self): + self.model = self.model_node.f( + "2", # kernel_size (string) + "2", # stride (string) + "0", # padding (string) + "1" # dilation (string) + )[0] + x = torch.ones(1, 3, 64, 64) + out = self.model(x) + self.assertEqual(out.size(), torch.Size([1, 3, 32, 32])) + + def test_kernel_stride_padding_as_tuple(self): + self.model = self.model_node.f( + "(2, 2)", # kernel_size (tuple as string) + "(2, 2)", # stride (tuple as string) + "(1, 1)", # padding (tuple as string) + "(1, 1)" # dilation (tuple as string) + )[0] + x = torch.ones(1, 3, 64, 64) + out = self.model(x) + self.assertEqual(out.size(), torch.Size([1, 3, 33, 33])) + + def test_padding_as_different_values(self): + for padding in ["0", "1"]: + with self.subTest(padding=padding): + self.model = self.model_node.f( + "2", # kernel_size (string) + "2", # stride (string) + padding, # padding (string) + "1" # dilation (string) + )[0] + x = torch.ones(1, 3, 64, 64) + out = self.model(x) + expected_size = 32 + int(padding) # Adjust for padding effect + self.assertEqual(out.size(), torch.Size([1, 3, expected_size, expected_size])) + + +if __name__ == "__main__": + unittest.main()