split subter implementation (training + inference)

Deep-SAD-PyTorch/src/networks/main.py

@@ -1,11 +1,12 @@
-from .mnist_LeNet import MNIST_LeNet, MNIST_LeNet_Autoencoder
-from .elpv_LeNet import ELPV_LeNet, ELPV_LeNet_Autoencoder
-from .subter_LeNet import SubTer_LeNet, SubTer_LeNet_Autoencoder
-from .fmnist_LeNet import FashionMNIST_LeNet, FashionMNIST_LeNet_Autoencoder
 from .cifar10_LeNet import CIFAR10_LeNet, CIFAR10_LeNet_Autoencoder
-from .mlp import MLP, MLP_Autoencoder
-from .vae import VariationalAutoencoder
 from .dgm import DeepGenerativeModel, StackedDeepGenerativeModel
+from .elpv_LeNet import ELPV_LeNet, ELPV_LeNet_Autoencoder
+from .fmnist_LeNet import FashionMNIST_LeNet, FashionMNIST_LeNet_Autoencoder
+from .mlp import MLP, MLP_Autoencoder
+from .mnist_LeNet import MNIST_LeNet, MNIST_LeNet_Autoencoder
+from .subter_LeNet import SubTer_LeNet, SubTer_LeNet_Autoencoder
+from .subter_LeNet_Split import SubTer_LeNet_Split, SubTer_LeNet_Split_Autoencoder
+from .vae import VariationalAutoencoder
 
 
 def build_network(net_name, ae_net=None):
@@ -15,6 +16,7 @@ def build_network(net_name, ae_net=None):
         "mnist_LeNet",
         "elpv_LeNet",
         "subter_LeNet",
+        "subter_LeNet_Split",
         "mnist_DGM_M2",
         "mnist_DGM_M1M2",
         "fmnist_LeNet",
@@ -46,6 +48,9 @@ def build_network(net_name, ae_net=None):
     if net_name == "subter_LeNet":
         net = SubTer_LeNet()
 
+    if net_name == "subter_LeNet_Split":
+        net = SubTer_LeNet_Split()
+
     if net_name == "elpv_LeNet":
         net = ELPV_LeNet()
 
@@ -130,6 +135,7 @@ def build_autoencoder(net_name):
     implemented_networks = (
         "elpv_LeNet",
         "subter_LeNet",
+        "subter_LeNet_Split",
         "mnist_LeNet",
         "mnist_DGM_M1M2",
         "fmnist_LeNet",
@@ -154,6 +160,9 @@ def build_autoencoder(net_name):
     if net_name == "subter_LeNet":
         ae_net = SubTer_LeNet_Autoencoder()
 
+    if net_name == "subter_LeNet_Split":
+        ae_net = SubTer_LeNet_Split_Autoencoder()
+
     if net_name == "elpv_LeNet":
         ae_net = ELPV_LeNet_Autoencoder()
 

Deep-SAD-PyTorch/src/networks/subter_LeNet_Split.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from base.base_net import BaseNet
+
+
+class SubTer_LeNet_Split(BaseNet):
+    def __init__(self, rep_dim=256):
+        super().__init__()
+
+        self.rep_dim = rep_dim
+        self.pool = nn.MaxPool2d(2, 2)
+
+        self.conv1 = nn.Conv2d(1, 8, 5, bias=False, padding=2)
+        self.bn1 = nn.BatchNorm2d(8, eps=1e-04, affine=False)
+        self.conv2 = nn.Conv2d(8, 4, 5, bias=False, padding=2)
+        self.bn2 = nn.BatchNorm2d(4, eps=1e-04, affine=False)
+        self.fc1 = nn.Linear(4 * 64 * 4, self.rep_dim, bias=False)
+
+    def forward(self, x):
+        x = x.view(-1, 1, 16, 256)
+        x = self.conv1(x)
+        x = self.pool(F.leaky_relu(self.bn1(x)))
+        x = self.conv2(x)
+        x = self.pool(F.leaky_relu(self.bn2(x)))
+        x = x.view(int(x.size(0)), -1)
+        x = self.fc1(x)
+        return x
+
+
+class SubTer_LeNet_Split_Decoder(BaseNet):
+    def __init__(self, rep_dim=256):
+        super().__init__()
+
+        self.rep_dim = rep_dim
+
+        # Decoder network
+        self.fc3 = nn.Linear(self.rep_dim, 4 * 64 * 4, bias=False)
+        self.bn3 = nn.BatchNorm2d(4, eps=1e-04, affine=False)
+        self.deconv1 = nn.ConvTranspose2d(4, 8, 5, bias=False, padding=2)
+        self.bn4 = nn.BatchNorm2d(8, eps=1e-04, affine=False)
+        self.deconv2 = nn.ConvTranspose2d(8, 1, 5, bias=False, padding=2)
+
+    def forward(self, x):
+        x = self.fc3(x)
+        x = x.view(int(x.size(0)), 4, 4, 64)
+        x = F.interpolate(F.leaky_relu(self.bn3(x)), scale_factor=2)
+        x = self.deconv1(x)
+        x = F.interpolate(F.leaky_relu(self.bn4(x)), scale_factor=2)
+        x = self.deconv2(x)
+        x = torch.sigmoid(x)
+        return x
+
+
+class SubTer_LeNet_Split_Autoencoder(BaseNet):
+    def __init__(self, rep_dim=256):
+        super().__init__()
+
+        self.rep_dim = rep_dim
+        self.encoder = SubTer_LeNet_Split(rep_dim=rep_dim)
+        self.decoder = SubTer_LeNet_Split_Decoder(rep_dim=rep_dim)
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.decoder(x)
+        return x