Extra arguments for instantiating convolutional layers

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>

Testing extra kwargs

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>

Author: Joao-L-S-Almeida

simulai/models/_pytorch_models/_autoencoder.py

@@ -1098,6 +1098,7 @@ def __init__(
         scale: float = 1e-3,
         devices: Union[str, list] = "cpu",
         name: str = None,
+        **kwargs,
     ) -> None:
         """
         Constructor method.
@@ -1172,6 +1173,7 @@ def __init__(
                 shallow=shallow,
                 use_batch_norm=use_batch_norm,
                 name=self.name,
+                **kwargs
             )
 
         self.encoder = encoder.to(self.device)

simulai/templates/_templates.py

@@ -673,6 +673,7 @@ def cnn_autoencoder_auto(
     use_batch_norm: bool = False,
     shallow: bool = False,
     name: str = None,
+    **kwargs,
 ) -> Tuple[NetworkTemplate, ...]:
 
     """
@@ -737,7 +738,7 @@ def cnn_autoencoder_auto(
 
     autogen_cnn = NetworkInstanceGen(
         architecture="cnn", dim=case, use_batch_norm=use_batch_norm,
-        kernel_size=kernel_size,
+        kernel_size=kernel_size, **kwargs,
     )
 
     autogen_dense = NetworkInstanceGen(architecture="dense", shallow=shallow)
@@ -798,6 +799,7 @@ def autoencoder_auto(
     use_batch_norm: bool = False,
     case: str = None,
     name: str = None,
+    **kwargs,
 ) -> Tuple[Union[NetworkTemplate, None], ...]:
 
     """
@@ -864,6 +866,7 @@ def autoencoder_auto(
             shallow=shallow,
             use_batch_norm=use_batch_norm,
             name=name,
+            **kwargs,
         )
 
         return encoder, decoder, bottleneck_encoder, bottleneck_decoder

tests/network/test_template_gen.py

@@ -372,6 +372,7 @@ def test_autoencoder_kernel_size_shallow(self) -> None:
             architecture="cnn",
             case="2d",
             shallow=True,
+            padding_mode='replicate',
         )
 
         estimated_data = autoencoder.eval(input_data=input_data)
@@ -392,6 +393,7 @@ def test_autoencoder_multiscaleautoencoder(self) -> None:
             case="2d",
             shallow=True,
             name="model",
+            padding_mode='replicate',
         )
 
         estimated_data = autoencoder.reconstruction_forward(input_data=input_data)