Bugfix/equiv padding weightedmse

cyto_dl/models/vae/base_vae.py

@@ -139,7 +139,10 @@ def __init__(
         self.latent_dim = latent_dim
 
         if decoder_latent_parts is None:
-            self.decoder_latent_parts = {key: self.prior.keys() for key in self.decoder.keys()}
+            pass
+            # self.decoder_latent_parts = {
+            #     key: self.prior.keys() for key in self.decoder.keys()
+            # }
         else:
             self.decoder_latent_parts = decoder_latent_parts
             for key in self.decoder.keys():
@@ -157,9 +160,11 @@ def __init__(
     def calculate_rcl(self, x, xhat, input_key, target_key=None):
         if not target_key:
             target_key = input_key
+
         rcl_per_input_dimension = self.reconstruction_loss[input_key](
             x[target_key], xhat[input_key]
         )
+
         return rcl_per_input_dimension
 
     def calculate_rcl_dict(self, x, xhat, z):
@@ -190,6 +195,10 @@ def calculate_elbo(self, x, xhat, z):
 
         total_kld = sum(kld_per_part_summed.values())
         total_recon = sum(rcl_reduced.values())
+        if len(total_recon.shape) > 0:
+            total_recon = total_recon.mean()
+            for key in rcl_reduced.keys():
+                rcl_reduced[key] = rcl_reduced[key].mean()
 
         return (
             total_recon + self.beta * total_kld,

cyto_dl/models/vae/image_encoder.py

@@ -123,6 +123,8 @@ def make_block(
                 subunits=self.num_res_units,
                 bias=bias,
             )
+        if padding is None:
+            padding = same_padding(kernel_size)
 
         return Convolution(
             spatial_dims=self.spatial_dims,

cyto_dl/models/vae/image_vae.py

@@ -34,6 +34,7 @@ def forward(self, x):
 class ImageVAE(BaseVAE):
     def __init__(
         self,
+        x_label: str,
         latent_dim: int,
         spatial_dims: int,
         in_shape: Sequence[int],
@@ -61,14 +62,15 @@ def __init__(
         num_res_units: int = 2,
         up_kernel_size: int = 3,
         first_conv_padding_mode: str = "replicate",
-        encoder_padding: Optional[Union[int, Sequence[int]]] = None,
         eps: float = 1e-8,
+        encoder_padding: Optional[Union[int, Sequence[int]]] = None,
+        metric_keys: Optional[list] = None,
         **base_kwargs,
     ):
         in_channels, *in_shape = in_shape
 
         self.out_channels = out_channels if out_channels is not None else in_channels
-
+        self.x_label = x_label
         self.spatial_dims = spatial_dims
         self.final_size = np.asarray(in_shape, dtype=int)
         self.up_kernel_size = up_kernel_size
@@ -207,6 +209,8 @@ def __init__(
             decoder=decoder,
             latent_dim=latent_dim,
             prior=prior,
+            x_label=x_label,
+            metric_keys=metric_keys,
             **base_kwargs,
         )