Fix up models

src/models/histogram_net.py

@@ -0,0 +1,150 @@
+import logging
+import torch
+import torch.nn as nn
+
+
+EPSILON = 1e-5
+
+
+class HistogramNet(nn.Module):
+    def __init__(
+        self,
+        elu_alpha: float = 1,
+        dropout_prob: float = 0.05,
+        features: list[int] = [16, 32, 64],
+        kernel_size: int = 3,
+        use_instance_norm: bool = True,
+        use_elu: bool = True,
+        leaky_relu_alpha: float = 0.01,
+        use_residual: bool = True,
+        **_,
+    ):
+        super(HistogramNet, self).__init__()
+        self._elu_alpha = elu_alpha
+        self._dropout_prob = dropout_prob
+        self._features = features
+        self._kernel_size = kernel_size
+        self._use_instance_norm = use_instance_norm
+        self._use_elu = use_elu
+        self._leaky_relu_alpha = leaky_relu_alpha
+        self._use_residual = use_residual
+        self.print_og_result = False
+
+        self._convolutions = nn.ModuleList(
+            self._make_conv_layer(in_channels=in_channels, out_channels=out_channels)
+            for in_channels, out_channels in zip([1] + features[:-1], features)
+        )
+
+        if self._use_residual:
+            self._residual_blocks = nn.ModuleList(
+                self._make_resblock(channels) for channels in features
+            )
+
+        self._deconvolutions = nn.ModuleList(
+            self._make_deconv_layer(in_channels=in_channels, out_channels=out_channels)
+            for in_channels, out_channels in zip(
+                features[::-1], features[::-1][1:] + [1]
+            )
+        )
+
+        self._initialize_weights()
+
+    def _make_conv_layer(self, in_channels: int, out_channels: int) -> nn.Sequential:
+        return nn.Sequential(
+            nn.Conv3d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=self._kernel_size,
+                padding=1,
+                bias=False,
+            ),
+            (
+                nn.ELU(self._elu_alpha)
+                if self._use_elu
+                else nn.LeakyReLU(self._leaky_relu_alpha)
+            ),
+            (nn.InstanceNorm3d if self._use_instance_norm else nn.BatchNorm3d)(
+                out_channels
+            ),
+            nn.Dropout(p=self._dropout_prob),
+        )
+
+    def _make_resblock(self, channels: int) -> nn.Sequential:
+        return nn.Sequential(
+            nn.Conv3d(
+                in_channels=channels,
+                out_channels=channels,
+                kernel_size=self._kernel_size,
+                padding=1,
+                bias=False,
+            ),
+            (
+                nn.ELU(self._elu_alpha)
+                if self._use_elu
+                else nn.LeakyReLU(self._leaky_relu_alpha)(
+                    nn.InstanceNorm3d if self._use_instance_norm else nn.BatchNorm3d
+                )(channels)
+            ),
+            nn.Conv3d(
+                in_channels=channels,
+                out_channels=channels,
+                kernel_size=self._kernel_size,
+                padding=1,
+                bias=False,
+            ),
+            (
+                nn.ELU(self._elu_alpha)
+                if self._use_elu
+                else nn.LeakyReLU(self._leaky_relu_alpha)
+            ),
+            (nn.InstanceNorm3d if self._use_instance_norm else nn.BatchNorm3d)(
+                channels
+            ),
+        )
+
+    def _make_deconv_layer(self, in_channels: int, out_channels: int) -> nn.Sequential:
+        return nn.Sequential(
+            nn.ConvTranspose3d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=self._kernel_size,
+                padding=1,
+            ),
+            (
+                nn.ELU(self._elu_alpha)
+                if self._use_elu
+                else nn.LeakyReLU(self._leaky_relu_alpha)
+            ),
+        )
+
+    def forward(self, x):
+        if self._use_residual:
+            for conv, res in zip(self._convolutions, self._residual_blocks):
+                x = conv(x)
+                x = x + res(x)
+        else:
+            for conv in self._convolutions:
+                x = conv(x)
+
+        for deconv in self._deconvolutions:
+            x = deconv(x)
+
+        if self.print_og_result:
+            logging.info(f"Original result {torch.sum(x)}")
+            self.print_og_result = False
+
+        return self._normalize(x)
+
+    @staticmethod
+    def _normalize(x):
+        x = torch.clamp(x, min=0)
+        x_sum = torch.sum(x, dim=(2, 3, 4), keepdim=True)
+        return x / (x_sum + EPSILON)
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, (nn.Conv3d, nn.ConvTranspose3d)):
+                # Applying He normal initialization
+                nn.init.kaiming_normal_(m.weight, mode="fan_in", nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)