Updates related to a unit test.

kaanaksit · Dec 19, 2023 · 7c63f55 · 7c63f55
1 parent 425cbe7
commit 7c63f55
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 15 deletions.
diff --git a/odak/learn/models/components.py b/odak/learn/models/components.py
@@ -683,7 +683,7 @@ def __init__(
                  no_spatial = False
                 ):
         """
-        Initializes the CBAM module.
+        Initializes the convolutional block attention module.
 
         Parameters
         ----------
@@ -713,7 +713,7 @@ def forward(self, x):
 
     def forward(self, x):
         """
-        Forward pass of the CBAM module.
+        Forward pass of the convolutional block attention module.
 
         Parameters
         ----------
@@ -746,8 +746,9 @@ def __init__(self, L):
                               Positional encoding level.
         """
         super(positional_encoder, self).__init__()
-        self.L=L
-
+        self.L = L
+
+
     def forward(self, x):
         """
         Forward model.
@@ -774,6 +775,4 @@ def forward(self, x):
         results = torch.cat(results, dim=2)
         results = results.permute(0, 2, 1)
         results = results.reshape(B, -1)
-        return results
-
-
+        return results 
diff --git a/odak/learn/models/models.py b/odak/learn/models/models.py
@@ -45,7 +45,7 @@ def __init__(self,
         self.siren_multiplier = siren_multiplier
         self.dimensions = dimensions
         for i in range(len(self.dimensions) - 1):
-            self.layers.append(torch.nn.Linear(dimensions[i], dimensions[i + 1], bias = self.bias))
+            self.layers.append(torch.nn.Linear(self.dimensions[i], self.dimensions[i + 1], bias = self.bias))
         if not isinstance(input_multiplier, type(None)):
             self.input_multiplier = torch.nn.ParameterList()
             self.input_multiplier.append(torch.nn.Parameter(torch.ones(1, self.dimensions[0]) * input_multiplier))

diff --git a/test/test_learn_models_positional_encoder.py b/test/test_learn_models_positional_encoder.py
@@ -7,19 +7,19 @@
 
 
 def main():
+    device_name = 'cpu'
     filename = './test/data/fruit_lady.png'
     test_filename  = './estimation.png'
     weights_filename = 'model_weights.pt'
     learning_rate = 1e-4
-    no_epochs = 25000
+    no_epochs = 5
+    save_at_every = 5000
     number_of_batches = 1
-    dimensions = [2, 256, 256, 256, 3]
     positional_encoding_level = 24
-    device_name = 'cuda'
-    save_at_every = 5000
-    device = torch.device(device_name)
-    positional_encoder = odak.learn.models.components.positional_encoder(L=positional_encoding_level)
+    dimensions = [2, 256, 256, 256, 3]
     dimensions[0] = dimensions[0] + dimensions[0] * 2 * positional_encoding_level
+    device = torch.device(device_name)
+    positional_encoder = odak.learn.models.components.positional_encoder(L = positional_encoding_level)
     model = odak.learn.models.multi_layer_perceptron(
                                                      dimensions = dimensions,
                                                      activation = torch.nn.ReLU(),
@@ -90,7 +90,6 @@ def trial(output_values, input_values, loss_function, model, positional_encoder)
         normalized_input_value[:, 1] = input_value[:, 1] / output_values.shape[1]
         input_value = positional_encoder(input_value)
         estimation = model(input_value)
-        estimation = model(input_value)
         ground_truth = output_values[input_value[:, 0].int(), input_value[:, 1].int(), :]
         estimated_image[input_value[:, 0].int(), input_value[:, 1].int(), :] = estimation
         loss = loss_function(estimation, ground_truth)