Issue baler-collaboration#305 - Block Datasets and GDN Activation Fun…

…ction (baler-collaboration#306)

* Issue 293 Imroved Conv3D Net

* Issue baler-collaboration#299 - Add Hurricane Isabel Dataset

* Issue baler-collaboration#301 - Converting Data into Smaller Blocks

* Adding Conv Blocking Logic

* Issue baler-collaboration#305 - Blocking of Data | GDN Activation Function

* Issue baler-collaboration#305 - Blocking | GDN Cleanup

* Issue baler-collaboration#305 - Blocking | GDN Cleanup 2

---------

Co-authored-by: Fritjof Bengtsson <[email protected]>

.gitignore

@@ -1,7 +1,7 @@
 /workspaces/*
 !/workspaces/public_datasets
-!workspaces/CMS_example/example_CMS/config/example_CMS_config.py
-!workspaces/CMS_example/example_CFD/config/example_CFD_config.py
+!workspaces/CMS_project_v1/CMS_project_v1/config/CMS_project_v1_config.py
+!workspaces/CFD_workspace/CFD_project_animation/config/CFD_project_animation_config.py
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

baler/baler.py

@@ -92,6 +92,7 @@ def perform_training(output_path, config, verbose: bool):
         config.custom_norm,
         config.test_size,
         config.apply_normalization,
+        config.convert_to_blocks,
     )
 
     if verbose:
@@ -318,6 +319,7 @@ def perform_decompression(output_path, config, verbose: bool):
         ),
         model_name=model_name,
         config=config,
+        output_path=output_path,
     )
     if verbose:
         print(f"Model used: {model_name}")

baler/modules/data_processing.py

@@ -23,6 +23,15 @@
 from ..modules import models
 
 
+def convert_to_blocks_util(blocks, data):
+    print("Converted Dataset to Blocks of Size - ", blocks)
+    blocks = np.array(blocks)
+    original_shape = np.array(data.shape)
+    total_size = np.prod(original_shape)
+    data = data.reshape((total_size // (blocks[1] * blocks[2])), blocks[1], blocks[2])
+    return data
+
+
 def save_model(model, model_path: str) -> None:
     """Saves the models state dictionary as a `.pt` file to the given path.
 

baler/modules/helper.py

@@ -265,7 +265,7 @@ def normalize(data, custom_norm):
     return data
 
 
-def process(input_path, custom_norm, test_size, apply_normalization):
+def process(input_path, custom_norm, test_size, apply_normalization, convert_to_blocks):
     """Loads the input data into a ndarray, splits it into train/test splits and normalizes if chosen.
 
     Args:
@@ -280,6 +280,9 @@ def process(input_path, custom_norm, test_size, apply_normalization):
     loaded = np.load(input_path)
     data = loaded["data"]
 
+    if convert_to_blocks:
+        data = data_processing.convert_to_blocks_util(convert_to_blocks, data)
+
     normalization_features = data_processing.find_minmax(data)
     if apply_normalization:
         print("Normalizing the data...")
@@ -453,6 +456,12 @@ def compress(model_path, config):
     # Loads the data and applies normalization if config.apply_normalization = True
     loaded = np.load(config.input_path)
     data_before = loaded["data"]
+
+    if config.convert_to_blocks:
+        data_before = data_processing.convert_to_blocks_util(
+            config.convert_to_blocks, data_before
+        )
+
     if config.apply_normalization:
         print("Normalizing...")
         data = normalize(data_before, config.custom_norm)
@@ -575,7 +584,13 @@ def compress(model_path, config):
 
 
 def decompress(
-    model_path, input_path, input_path_deltas, input_batch_index, model_name, config
+    model_path,
+    input_path,
+    input_path_deltas,
+    input_batch_index,
+    model_name,
+    config,
+    output_path,
 ):
     """Function which performs the decompression of the compressed file. In order to decompress, you must have a
     compressed file, whose path is determined by `input_path`, a model from path `model_path` and a model_name. The
@@ -598,6 +613,11 @@ def decompress(
     names = loaded["names"]
     normalization_features = loaded["normalization_features"]
 
+    if config.model_type == "convolutional":
+        final_layer_details = np.load(
+            os.path.join(output_path, "training", "final_layer.npy"),
+        )
+
     if config.save_error_bounded_deltas:
         loaded_deltas = np.load(
             gzip.GzipFile(input_path_deltas, "r"), allow_pickle=True
@@ -630,6 +650,9 @@ def decompress(
     )
     model.eval()
 
+    if config.model_type == "convolutional":
+        model.set_final_layer_dims(final_layer_details)
+
     # Load the data, convert to tensor and batch it to avoid memory leaks
     data_tensor = torch.from_numpy(data).to(device)
     data_dl = DataLoader(

baler/modules/models.py

@@ -17,6 +17,102 @@
 from torch.nn import functional as F
 
 
+import torch.utils.data
+from torch.nn import functional as F
+from torch.autograd import Function
+from ..modules import helper
+
+
+class LowerBound(Function):
+    @staticmethod
+    def forward(ctx, inputs, bound):
+        b = torch.ones(inputs.size(), device=inputs.device) * bound
+        b = b.to(inputs.device)
+        ctx.save_for_backward(inputs, b)
+        return torch.max(inputs, b)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        inputs, b = ctx.saved_tensors
+
+        pass_through_1 = inputs >= b
+        pass_through_2 = grad_output < 0
+
+        pass_through = pass_through_1 | pass_through_2
+        return pass_through.type(grad_output.dtype) * grad_output, None
+
+
+class GDN(nn.Module):
+    """Generalized divisive normalization layer.
+    y[i] = x[i] / sqrt(beta[i] + sum_j(gamma[j, i] * x[j]^2))
+    """
+
+    def __init__(
+        self,
+        ch,
+        inverse=False,
+        beta_min=1e-6,
+        gamma_init=0.1,
+        reparam_offset=2**-18,
+    ):
+        super(GDN, self).__init__()
+        self.inverse = inverse
+        self.beta_min = beta_min
+        self.gamma_init = gamma_init
+        self.device = helper.get_device()
+        self.reparam_offset = torch.tensor([reparam_offset], device=self.device)
+
+        self.build(ch)
+
+    def build(self, ch):
+        self.pedestal = self.reparam_offset**2
+        self.beta_bound = (self.beta_min + self.reparam_offset**2) ** 0.5
+        self.gamma_bound = self.reparam_offset
+
+        # Create beta param
+        beta = torch.sqrt(torch.ones(ch, device=self.device) + self.pedestal)
+        self.beta = nn.Parameter(beta)
+
+        # Create gamma param
+        eye = torch.eye(ch, device=self.device)
+        g = self.gamma_init * eye
+        g = g + self.pedestal
+        gamma = torch.sqrt(g)
+        self.gamma = nn.Parameter(gamma)
+
+    def forward(self, inputs):
+        unfold = False
+        if inputs.dim() == 5:
+            unfold = True
+            bs, ch, d, w, h = inputs.size()
+            inputs = inputs.view(bs, ch, d * w, h)
+
+        _, ch, _, _ = inputs.size()
+
+        # Beta bound and reparam
+        beta = LowerBound.apply(self.beta, self.beta_bound)
+        beta = beta**2 - self.pedestal
+
+        # Gamma bound and reparam
+        gamma = LowerBound.apply(self.gamma, self.gamma_bound)
+        gamma = gamma**2 - self.pedestal
+        gamma = gamma.view(ch, ch, 1, 1)
+
+        # Norm pool calc
+        norm_ = nn.functional.conv2d(inputs**2, gamma, beta)
+        norm_ = torch.sqrt(norm_)
+
+        # Apply norm
+        if self.inverse:
+            outputs = inputs * norm_
+        else:
+            outputs = inputs / norm_
+
+        if unfold:
+            outputs = outputs.view(bs, ch, d, w, h)
+        return outputs
+
+
 class AE(nn.Module):
     # This class is a modified version of the original class by George Dialektakis found at
     # https://github.com/Autoencoders-compression-anomaly/Deep-Autoencoders-Data-Compression-GSoC-2021
@@ -224,7 +320,8 @@ def __init__(self, n_features, z_dim, *args, **kwargs):
         super(Conv_AE, self).__init__(*args, **kwargs)
 
         self.q_z_mid_dim = 2000
-        self.q_z_output_dim = 72128
+        self.q_z_output_dim = 128
+        self.conv_op_shape = None
 
         # Encoder
 
@@ -242,6 +339,7 @@ def __init__(self, n_features, z_dim, *args, **kwargs):
         )
         # Flatten
         self.flatten = nn.Flatten(start_dim=1)
+
         # Linear layers
         self.q_z_lin = nn.Sequential(
             nn.Linear(self.q_z_output_dim, self.q_z_mid_dim),
@@ -276,6 +374,9 @@ def __init__(self, n_features, z_dim, *args, **kwargs):
     def encode(self, x):
         # Conv
         out = self.q_z_conv(x)
+        self.conv_op_shape = out.shape
+        self.q_z_output_dim = out.shape[1] * out.shape[2] * out.shape[3]
+
         # Flatten
         out = self.flatten(out)
         # Dense
@@ -286,7 +387,12 @@ def decode(self, z):
         # Dense
         out = self.p_x_lin(z)
         # Unflatten
-        out = out.view(out.size(0), 32, 49, 46)
+        out = out.view(
+            self.conv_op_shape[0],
+            self.conv_op_shape[1],
+            self.conv_op_shape[2],
+            self.conv_op_shape[3],
+        )
         # Conv transpose
         out = self.p_x_conv(out)
         return out
@@ -296,6 +402,12 @@ def forward(self, x):
         out = self.decode(z)
         return out
 
+    def get_final_layer_dims(self):
+        return self.conv_op_shape
+
+    def set_final_layer_dims(self, conv_op_shape):
+        self.conv_op_shape = conv_op_shape
+
 
 class FPGA_prototype_model(nn.Module):
     def __init__(self, n_features, z_dim, *args, **kwargs):
@@ -459,3 +571,97 @@ def forward(self, x):
 
     def set_compress_to_latent_space(self, compress_to_latent_space):
         self.compress_to_latent_space = compress_to_latent_space
+
+
+class Conv_AE_GDN(nn.Module):
+    def __init__(self, n_features, z_dim, *args, **kwargs):
+        super(Conv_AE_GDN, self).__init__(*args, **kwargs)
+
+        self.q_z_mid_dim = 2000
+        self.q_z_output_dim = 128
+        self.conv_op_shape = None
+
+        # Encoder
+
+        # Conv Layers
+        self.q_z_conv = nn.Sequential(
+            nn.Conv2d(1, 8, kernel_size=(2, 5), stride=(1), padding=(1)),
+            # nn.BatchNorm2d(8),
+            GDN(8),
+            nn.Conv2d(8, 16, kernel_size=(3), stride=(1), padding=(1)),
+            nn.BatchNorm2d(16),
+            GDN(16),
+            nn.Conv2d(16, 32, kernel_size=(3), stride=(1), padding=(0)),
+            # nn.BatchNorm2d(32),
+            GDN(32),
+        )
+        # Flatten
+        self.flatten = nn.Flatten(start_dim=1)
+
+        # Linear layers
+        self.q_z_lin = nn.Sequential(
+            nn.Linear(self.q_z_output_dim, self.q_z_mid_dim),
+            nn.ReLU(),
+            # nn.BatchNorm1d(self.q_z_output_dim),
+            nn.Linear(self.q_z_mid_dim, z_dim),
+            nn.ReLU(),
+        )
+
+        # Decoder
+
+        # Linear layers
+        self.p_x_lin = nn.Sequential(
+            nn.Linear(z_dim, self.q_z_mid_dim),
+            nn.ReLU(),
+            # nn.BatchNorm1d(self.q_z_output_dim),
+            nn.Linear(self.q_z_mid_dim, self.q_z_output_dim),
+            nn.ReLU()
+            # nn.BatchNorm1d(42720)
+        )
+        # Conv Layers
+        self.p_x_conv = nn.Sequential(
+            nn.ConvTranspose2d(32, 16, kernel_size=(3), stride=(1), padding=(0)),
+            # nn.BatchNorm2d(16),
+            GDN(16, inverse=True),
+            nn.ConvTranspose2d(16, 8, kernel_size=(3), stride=(1), padding=(1)),
+            # nn.BatchNorm2d(8),
+            GDN(8, inverse=True),
+            nn.ConvTranspose2d(8, 1, kernel_size=(2, 5), stride=(1), padding=(1)),
+        )
+
+    def encode(self, x):
+        # Conv
+        out = self.q_z_conv(x)
+        self.conv_op_shape = out.shape
+        self.q_z_output_dim = out.shape[1] * out.shape[2] * out.shape[3]
+
+        # Flatten
+        out = self.flatten(out)
+        # Dense
+        out = self.q_z_lin(out)
+        return out
+
+    def decode(self, z):
+        # Dense
+        out = self.p_x_lin(z)
+        # Unflatten
+        out = out.view(
+            self.conv_op_shape[0],
+            self.conv_op_shape[1],
+            self.conv_op_shape[2],
+            self.conv_op_shape[3],
+        )
+        # Conv transpose
+        out = self.p_x_conv(out)
+        return out
+
+    def forward(self, x):
+        z = self.encode(x)
+        out = self.decode(z)
+        return out
+
+    def get_final_layer_dims(self):
+        return self.conv_op_shape
+
+    def set_final_layer_dims(self, conv_op_shape):
+        self.conv_op_shape = conv_op_shape