Update networks

src/editor/models/residual3.py

@@ -0,0 +1,146 @@
+import torch
+import torch.nn as nn
+
+
+class DepthwiseSeparableConv3d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, padding):
+        super(DepthwiseSeparableConv3d, self).__init__()
+        self.depthwise = nn.Conv3d(
+            in_channels,
+            in_channels,
+            kernel_size=kernel_size,
+            padding=padding,
+            groups=in_channels,
+        )
+        self.pointwise = nn.Conv3d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        x = self.depthwise(x)
+        x = self.pointwise(x)
+        return x
+
+
+class Residual3(nn.Module):
+    def __init__(
+        self,
+        elu_alpha: float = 1,
+        dropout_prob: float = 0.1,
+        use_depthwise_separable_conv: bool = False,
+        feature_map_sizes: list[int] = [16, 32, 64],
+        kernel_sizes: list[int] = [3, 3, 3],
+    ):
+        super(Residual3, self).__init__()
+
+        conv = DepthwiseSeparableConv3d if use_depthwise_separable_conv else nn.Conv3d
+
+        # Assuming the input histograms are 3D tensors of shape (bin_count, bin_count, bin_count)
+        # Convolutional layers to extract features from the histograms
+        self.conv1 = conv(
+            1, feature_map_sizes[0], kernel_size=kernel_sizes[0], padding=1
+        )
+        self.conv2 = conv(
+            feature_map_sizes[0],
+            feature_map_sizes[1],
+            kernel_size=kernel_sizes[1],
+            padding=1,
+        )
+        self.conv3 = conv(
+            feature_map_sizes[1],
+            feature_map_sizes[2],
+            kernel_size=kernel_sizes[2],
+            padding=1,
+        )
+
+        self.activation = nn.ELU(elu_alpha, inplace=True)
+
+        self.bn1 = nn.BatchNorm3d(feature_map_sizes[0])
+        self.bn2 = nn.BatchNorm3d(feature_map_sizes[1])
+        self.bn3 = nn.BatchNorm3d(feature_map_sizes[2])
+
+        self.resblock1 = nn.Sequential(
+            conv(
+                feature_map_sizes[2],
+                feature_map_sizes[2],
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=False,
+            ),
+            nn.ELU(elu_alpha, inplace=True),
+            nn.BatchNorm3d(feature_map_sizes[2]),
+            conv(
+                feature_map_sizes[2],
+                feature_map_sizes[2],
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=False,
+            ),
+            nn.ELU(elu_alpha, inplace=True),
+            nn.BatchNorm3d(feature_map_sizes[2]),
+        )
+
+        # Deconvolutional layers
+        self.deconv1 = nn.ConvTranspose3d(
+            feature_map_sizes[2],
+            feature_map_sizes[1],
+            kernel_size=feature_map_sizes[2],
+            stride=1,
+            padding=1,
+        )
+        self.deconv2 = nn.ConvTranspose3d(
+            feature_map_sizes[1],
+            feature_map_sizes[0],
+            kernel_size=feature_map_sizes[1],
+            stride=1,
+            padding=1,
+        )
+        self.deconv3 = nn.ConvTranspose3d(
+            feature_map_sizes[0], 1, kernel_size=3, stride=1, padding=1
+        )
+
+        self.dropout = nn.Dropout3d(p=dropout_prob)
+        self._initialize_weights()
+
+    def forward(self, x):
+        out = self.dropout(self.bn1(self.activation(self.conv1(x))))
+        out = self.dropout(self.bn2(self.activation(self.conv2(out))))
+        out = self.dropout(self.bn3(self.activation(self.conv2(out))))
+
+        out = out + self.resblock1(out)
+
+        out = self.activation(self.deconv1(out))
+        out = self.activation(self.deconv2(out))
+        out = self.activation(self.deconv3(out))
+
+        return self._normalize(out)
+
+    def _normalize(self, x):
+        x_sum = torch.sum(x, dim=(2, 3, 4), keepdim=True)
+        return x / torch.where(x_sum == 0, torch.ones_like(x_sum), x_sum)
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv3d) or isinstance(m, nn.ConvTranspose3d):
+                nn.init.xavier_normal_(
+                    m.weight,
+                )
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+def _test_network_dimensions(constructor):
+    for bin_count in [16, 32, 64]:
+        model = constructor()
+
+        # Create a dummy input tensor of the correct shape
+        input_tensor = torch.rand(4, 1, bin_count, bin_count, bin_count)
+
+        # Test the model output
+        output = model(input_tensor)
+        assert (
+            input_tensor.shape == output.shape
+        ), f"Expected output shape {input_tensor.shape}, but got {output.shape}"
+
+
+_test_network_dimensions(Residual3)