nodes/HowToSD-ComfyUI-Pt-Wrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added Ptn Conv2d, Ptn Max Pool2d, Ptn Avg Pool2d nodes.

Browse Source
This commit is contained in:
HowToSD
2025-03-15 21:26:57 -07:00
parent 7197708ee9
commit f36a829fc6
23 changed files with 586 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2 -1
View File
@@ -3,4 +3,5 @@ __pycache__/
models/
*.pth
/datasets/
examples/datasets/.~lock*
examples/datasets/.~lock*
.vscode
__init__.py
+9
View File
@@ -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",
docs/reference/node_reference.md
+11 -5
View File
@@ -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. |
docs/reference/ptn_avg_pool2d.md
+24
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Model
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_chained_model.md
+1 -1
View File
@@ -19,6 +19,6 @@ Constructs a chained PyTorch model.
| Ptmodel |
<HR>
Category: PyTorch wrapper - Training
Category: PyTorch wrapper - Model
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_conv2d.md
+38
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Model
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_conv_model.md
+1 -1
View File
@@ -40,6 +40,6 @@ A convolutional model consisting of multiple convolutional layers.
| Ptmodel |
<HR>
Category: PyTorch wrapper - Training
Category: PyTorch wrapper - Model
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_linear.md
+1 -1
View File
@@ -19,6 +19,6 @@ A linear model consisting of a single dense layer.
| Ptmodel |
<HR>
Category: PyTorch wrapper - Training
Category: PyTorch wrapper - Model
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_linear_model.md
+1 -1
View File
@@ -30,6 +30,6 @@ A linear model consisting of dense layers.
| Ptmodel |
<HR>
Category: PyTorch wrapper - Training
Category: PyTorch wrapper - Model
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_max_pool2d.md
+26
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Model
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_resnet_model.md
+1 -1
View File
@@ -23,6 +23,6 @@ A Resnet model consisting of multiple Resnet layers.
| Ptmodel |
<HR>
Category: PyTorch wrapper - Training
Category: PyTorch wrapper - Model
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
modules/pytorch_wrapper/ptn_avg_pool2d.py
+68
View File
@@ -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,)
modules/pytorch_wrapper/ptn_chained_model.py
+1 -1
View File
@@ -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
modules/pytorch_wrapper/ptn_conv2d.py
+109
View File
@@ -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,)
modules/pytorch_wrapper/ptn_conv_model.py
+1 -1
View File
@@ -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
modules/pytorch_wrapper/ptn_linear.py
+1 -1
View File
@@ -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
modules/pytorch_wrapper/ptn_linear_model.py
+3 -3
View File
@@ -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}),
}
}
modules/pytorch_wrapper/ptn_max_pool2d.py
+75
View File
@@ -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,)
modules/pytorch_wrapper/ptn_resnet_model.py
+1 -1
View File
@@ -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
pyproject.toml
+1 -1
View File
@@ -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]
tests/pytorch_wrapper/ptn_avg_pool2d_test.py
+52
View File
@@ -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()
tests/pytorch_wrapper/ptn_conv2d_test.py
+104
View File
@@ -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()
tests/pytorch_wrapper/ptn_max_pool2d_test.py
+55
View File
@@ -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()