51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
|
import torch
|
||
|
|
from NNMF2dConvGroupedAutograd import NNMF2dConvGrouped
|
||
|
|
from append_parameter import append_parameter
|
||
|
|
|
||
|
|
|
||
|
|
def append_block(
|
||
|
|
network: torch.nn.Sequential,
|
||
|
|
out_channels: int,
|
||
|
|
test_image: torch.Tensor,
|
||
|
|
parameter_cnn_top: list[torch.nn.parameter.Parameter],
|
||
|
|
parameter_nnmf: list[torch.nn.parameter.Parameter],
|
||
|
|
parameter_norm: list[torch.nn.parameter.Parameter],
|
||
|
|
torch_device: torch.device,
|
||
|
|
dilation: tuple[int, int] | int = 1,
|
||
|
|
padding: tuple[int, int] | int = 0,
|
||
|
|
stride: tuple[int, int] | int = 1,
|
||
|
|
kernel_size: tuple[int, int] = (5, 5),
|
||
|
|
epsilon: float | None = None,
|
||
|
|
positive_function_type: int = 0,
|
||
|
|
beta: float | None = None,
|
||
|
|
iterations: int = 20,
|
||
|
|
local_learning: bool = False,
|
||
|
|
local_learning_kl: bool = False,
|
||
|
|
momentum: float = 0.1,
|
||
|
|
track_running_stats: bool = False,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
|
||
|
|
kernel_size_internal: list[int] = [kernel_size[-2], kernel_size[-1]]
|
||
|
|
|
||
|
|
if kernel_size[0] < 1:
|
||
|
|
kernel_size_internal[0] = test_image.shape[-2]
|
||
|
|
|
||
|
|
if kernel_size[1] < 1:
|
||
|
|
kernel_size_internal[1] = test_image.shape[-1]
|
||
|
|
|
||
|
|
network.append(
|
||
|
|
NNMF2dConvGrouped(
|
||
|
|
in_channels=test_image.shape[1],
|
||
|
|
out_channels=out_channels,
|
||
|
|
kernel_size=(kernel_size_internal[-2], kernel_size_internal[-1]),
|
||
|
|
dilation=dilation,
|
||
|
|
padding=padding,
|
||
|
|
stride=stride,
|
||
|
|
device=torch_device,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
test_image = network[-1](test_image)
|
||
|
|
append_parameter(module=network[-1], parameter_list=parameter_nnmf)
|
||
|
|
|
||
|
|
return test_image
|