Bernstein_Poster_2024/basis_nnmf_sp1.01_x4/make_network.py
David Rotermund 12eb04e446 New Sims
2024-11-05 18:20:02 +01:00

212 lines
6.1 KiB
Python

import torch
from append_block import append_block
from L1NormLayer import L1NormLayer
from NNMF2d import NNMF2d
from append_parameter import append_parameter
def make_network(
input_dim_x: int,
input_dim_y: int,
input_number_of_channel: int,
iterations: int,
torch_device: torch.device,
epsilon: bool | None = None,
positive_function_type: int = 0,
beta: float | None = None,
# Conv:
number_of_output_channels: list[int] = [32*4, 64*4, 96*4, 10],
kernel_size_conv: list[tuple[int, int]] = [
(5, 5),
(5, 5),
(-1, -1), # Take the whole input image x and y size
(1, 1),
],
stride_conv: list[tuple[int, int]] = [
(1, 1),
(1, 1),
(1, 1),
(1, 1),
],
padding_conv: list[tuple[int, int]] = [
(0, 0),
(0, 0),
(0, 0),
(0, 0),
],
dilation_conv: list[tuple[int, int]] = [
(1, 1),
(1, 1),
(1, 1),
(1, 1),
],
# Pool:
kernel_size_pool: list[tuple[int, int]] = [
(2, 2),
(2, 2),
(-1, -1), # No pooling layer
(-1, -1), # No pooling layer
],
stride_pool: list[tuple[int, int]] = [
(2, 2),
(2, 2),
(-1, -1),
(-1, -1),
],
padding_pool: list[tuple[int, int]] = [
(0, 0),
(0, 0),
(0, 0),
(0, 0),
],
dilation_pool: list[tuple[int, int]] = [
(1, 1),
(1, 1),
(1, 1),
(1, 1),
],
enable_onoff: bool = False,
) -> tuple[
torch.nn.Sequential,
list[list[torch.nn.parameter.Parameter]],
list[str],
]:
assert len(number_of_output_channels) == len(kernel_size_conv)
assert len(number_of_output_channels) == len(stride_conv)
assert len(number_of_output_channels) == len(padding_conv)
assert len(number_of_output_channels) == len(dilation_conv)
assert len(number_of_output_channels) == len(kernel_size_pool)
assert len(number_of_output_channels) == len(stride_pool)
assert len(number_of_output_channels) == len(padding_pool)
assert len(number_of_output_channels) == len(dilation_pool)
if enable_onoff:
input_number_of_channel *= 2
parameter_cnn_top: list[torch.nn.parameter.Parameter] = []
parameter_nnmf: list[torch.nn.parameter.Parameter] = []
parameter_norm: list[torch.nn.parameter.Parameter] = []
test_image = torch.ones(
(1, input_number_of_channel, input_dim_x, input_dim_y), device=torch_device
)
network = torch.nn.Sequential()
network = network.to(torch_device)
for block_id in range(0, len(number_of_output_channels)):
test_image = append_block(
network=network,
out_channels=number_of_output_channels[block_id],
test_image=test_image,
dilation=dilation_conv[block_id],
padding=padding_conv[block_id],
stride=stride_conv[block_id],
kernel_size=kernel_size_conv[block_id],
epsilon=epsilon,
positive_function_type=positive_function_type,
beta=beta,
iterations=iterations,
torch_device=torch_device,
parameter_cnn_top=parameter_cnn_top,
parameter_nnmf=parameter_nnmf,
parameter_norm=parameter_norm,
)
if (kernel_size_pool[block_id][0] > 0) and (kernel_size_pool[block_id][1] > 0):
network.append(torch.nn.ReLU())
test_image = network[-1](test_image)
mock_output = (
torch.nn.functional.conv2d(
torch.zeros(
1,
1,
test_image.shape[2],
test_image.shape[3],
),
torch.zeros((1, 1, 2, 2)),
stride=(2, 2),
padding=(0, 0),
dilation=(1, 1),
)
.squeeze(0)
.squeeze(0)
)
network.append(
torch.nn.Unfold(
kernel_size=(2, 2),
stride=(2, 2),
padding=(0, 0),
dilation=(1, 1),
)
)
test_image = network[-1](test_image)
network.append(
torch.nn.Fold(
output_size=mock_output.shape,
kernel_size=(1, 1),
dilation=1,
padding=0,
stride=1,
)
)
test_image = network[-1](test_image)
network.append(L1NormLayer())
test_image = network[-1](test_image)
network.append(
NNMF2d(
in_channels=test_image.shape[1],
out_channels=test_image.shape[1] // 4,
epsilon=epsilon,
positive_function_type=positive_function_type,
beta=beta,
iterations=iterations,
local_learning=False,
local_learning_kl=False,
).to(torch_device)
)
test_image = network[-1](test_image)
append_parameter(module=network[-1], parameter_list=parameter_nnmf)
network.append(
torch.nn.BatchNorm2d(
num_features=test_image.shape[1],
device=torch_device,
momentum=0.1,
track_running_stats=False,
)
)
test_image = network[-1](test_image)
append_parameter(module=network[-1], parameter_list=parameter_norm)
network.append(torch.nn.Softmax(dim=1))
test_image = network[-1](test_image)
network.append(torch.nn.Flatten())
test_image = network[-1](test_image)
parameters: list[list[torch.nn.parameter.Parameter]] = [
parameter_cnn_top,
parameter_nnmf,
parameter_norm,
]
name_list: list[str] = [
"cnn_top",
"nnmf",
"batchnorm2d",
]
return (
network,
parameters,
name_list,
)