careamics 0.0.4.2__py3-none-any.whl → 0.0.5__py3-none-any.whl

careamics/models/lvae/stochastic.py

@@ -0,0 +1,393 @@
+"""Script containing the common basic blocks (nn.Module)
+reused by the LadderVAE architecture.
+"""
+
+from typing import Dict, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as F
+from torch.distributions import kl_divergence
+from torch.distributions.normal import Normal
+
+from .utils import (
+    StableLogVar,
+    StableMean,
+    kl_normal_mc,
+)
+
+ConvType = Union[nn.Conv2d, nn.Conv3d]
+NormType = Union[nn.BatchNorm2d, nn.BatchNorm3d]
+DropoutType = Union[nn.Dropout2d, nn.Dropout3d]
+
+
+class NormalStochasticBlock(nn.Module):
+    """
+    Stochastic block used in the Top-Down inference pass.
+
+    Algorithm:
+        - map input parameters to q(z) and (optionally) p(z) via convolution
+        - sample a latent tensor z ~ q(z)
+        - feed z to convolution and return.
+
+    NOTE 1:
+        If parameters for q are not given, sampling is done from p(z).
+
+    NOTE 2:
+        The restricted KL divergence is obtained by first computing the element-wise KL divergence
+        (i.e., the KL computed for each element of the latent tensors). Then, the restricted version
+        is computed by summing over the channels and the spatial dimensions associated only to the
+        portion of the latent tensor that is used for prediction.
+    """
+
+    def __init__(
+        self,
+        c_in: int,
+        c_vars: int,
+        c_out: int,
+        conv_dims: int = 2,
+        kernel: int = 3,
+        transform_p_params: bool = True,
+        vanilla_latent_hw: int = None,
+        use_naive_exponential: bool = False,
+    ):
+        """
+        Parameters
+        ----------
+        c_in: int
+            The number of channels of the input tensor.
+        c_vars: int
+            The number of channels of the latent space tensor.
+        c_out:  int
+            The output of the stochastic layer.
+            Note that this is different from the sampled latent z.
+        conv_dims: int, optional
+            The number of dimensions of the convolutional layers (2D or 3D).
+            Default is 2.
+        kernel: int, optional
+            The size of the kernel used in convolutional layers.
+            Default is 3.
+        transform_p_params: bool, optional
+            Whether a transformation should be applied to the `p_params` tensor.
+            The transformation consists in a 2D convolution ()`conv_in_p()`) that
+            maps the input to a larger number of channels.
+            Default is `True`.
+        vanilla_latent_hw: int, optional
+            The shape of the latent tensor used for prediction (i.e., it influences the computation of restricted KL).
+            Default is `None`.
+        use_naive_exponential: bool, optional
+            If `False`, exponentials are computed according to the alternative definition
+            provided by `StableExponential` class. This should improve numerical stability
+            in the training process. Default is `False`.
+        """
+        super().__init__()
+        assert kernel % 2 == 1
+        pad = kernel // 2
+        self.transform_p_params = transform_p_params
+        self.c_in = c_in
+        self.c_out = c_out
+        self.c_vars = c_vars
+        self.conv_dims = conv_dims
+        self._use_naive_exponential = use_naive_exponential
+        self._vanilla_latent_hw = vanilla_latent_hw
+
+        conv_layer: ConvType = getattr(nn, f"Conv{conv_dims}d")
+
+        if transform_p_params:
+            self.conv_in_p = conv_layer(c_in, 2 * c_vars, kernel, padding=pad)
+        self.conv_in_q = conv_layer(c_in, 2 * c_vars, kernel, padding=pad)
+        self.conv_out = conv_layer(c_vars, c_out, kernel, padding=pad)
+
+    def get_z(
+        self,
+        sampling_distrib: torch.distributions.normal.Normal,
+        forced_latent: Union[torch.Tensor, None],
+        mode_pred: bool,
+        use_uncond_mode: bool,
+    ) -> torch.Tensor:
+        """Sample a latent tensor from the given latent distribution.
+
+        Latent tensor can be obtained is several ways:
+            - Sampled from the (Gaussian) latent distribution.
+            - Taken as a pre-defined forced latent.
+            - Taken as the mode (mean) of the latent distribution.
+            - In prediction mode (`mode_pred==True`), can be either sample or taken as the distribution mode.
+
+        Parameters
+        ----------
+        sampling_distrib: torch.distributions.normal.Normal
+            The Gaussian distribution from which latent tensor is sampled.
+        forced_latent: torch.Tensor
+            A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent tensor and,
+            hence, sampling does not happen.
+        mode_pred: bool
+            Whether the model is prediction mode.
+        use_uncond_mode: bool
+            Whether to use the uncoditional distribution p(z) to sample latents in prediction mode.
+        """
+        if forced_latent is None:
+            if mode_pred:
+                if use_uncond_mode:
+                    z = sampling_distrib.mean
+                else:
+                    z = sampling_distrib.rsample()
+            else:
+                z = sampling_distrib.rsample()
+        else:
+            z = forced_latent
+        return z
+
+    def sample_from_q(
+        self, q_params: torch.Tensor, var_clip_max: float
+    ) -> torch.Tensor:
+        """
+        Given an input parameter tensor defining q(z),
+        it processes it by calling `process_q_params()` method and
+        sample a latent tensor from the resulting distribution.
+
+        Parameters
+        ----------
+        q_params: torch.Tensor
+            The input tensor to be processed.
+        var_clip_max: float
+            The maximum value reachable by the log-variance of the latent distribution.
+            Values exceeding this threshold are clipped.
+        """
+        _, _, q = self.process_q_params(q_params, var_clip_max)
+        return q.rsample()
+
+    def compute_kl_metrics(
+        self,
+        p: torch.distributions.normal.Normal,
+        p_params: torch.Tensor,
+        q: torch.distributions.normal.Normal,
+        q_params: torch.Tensor,
+        mode_pred: bool,
+        analytical_kl: bool,
+        z: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Compute KL (analytical or MC estimate) and then process it, extracting composed versions of the metric.
+        Specifically, the different versions of the KL loss terms are:
+            - `kl_elementwise`: KL term for each single element of the latent tensor [Shape: (batch, ch, h, w)].
+            - `kl_samplewise`: KL term associated to each sample in the batch [Shape: (batch, )].
+            - `kl_samplewise_restricted`: KL term only associated to the portion of the latent tensor that is
+            used for prediction and summed over channel and spatial dimensions [Shape: (batch, )].
+            - `kl_channelwise`: KL term associated to each sample and each channel [Shape: (batch, ch, )].
+            - `kl_spatial`: KL term summed over the channels, i.e., retaining the spatial dimensions [Shape: (batch, h, w)]
+
+        Parameters
+        ----------
+        p: torch.distributions.normal.Normal
+            The prior generative distribution p(z_i|z_{i+1}) (or p(z_L)).
+        p_params: torch.Tensor
+            The parameters of the prior generative distribution.
+        q: torch.distributions.normal.Normal
+            The inference distribution q(z_i|z_{i+1}) (or q(z_L|x)).
+        q_params: torch.Tensor
+            The parameters of the inference distribution.
+        mode_pred: bool
+            Whether the model is in prediction mode.
+        analytical_kl: bool
+            Whether to compute the KL divergence analytically or using Monte Carlo estimation.
+        z: torch.Tensor
+            The sampled latent tensor.
+        """
+        if mode_pred is False:  # if not predicting
+            if analytical_kl:
+                kl_elementwise = kl_divergence(q, p)
+            else:
+                kl_elementwise = kl_normal_mc(z, p_params, q_params)
+
+            all_dims = tuple(range(len(kl_elementwise.shape)))
+            kl_samplewise = kl_elementwise.sum(all_dims[1:])
+            kl_channelwise = kl_elementwise.sum(all_dims[2:])
+
+            # compute KL only on the portion of the latent space that is used for prediction.
+            pad = (kl_elementwise.shape[-1] - self._vanilla_latent_hw) // 2
+            if pad > 0:
+                tmp = kl_elementwise[..., pad:-pad, pad:-pad]
+                kl_samplewise_restricted = tmp.sum(all_dims[1:])
+            else:
+                kl_samplewise_restricted = kl_samplewise
+
+            # Compute spatial KL analytically (but conditioned on samples from
+            # previous layers)
+            kl_spatial = kl_elementwise.sum(1)
+        else:  # if predicting, no need to compute KL
+            kl_elementwise = kl_samplewise = kl_spatial = kl_channelwise = None
+
+        kl_dict = {
+            "kl_elementwise": kl_elementwise,  # (batch, ch, h, w)
+            "kl_samplewise": kl_samplewise,  # (batch, )
+            "kl_samplewise_restricted": kl_samplewise_restricted,  # (batch, )
+            "kl_spatial": kl_spatial,  # (batch, h, w)
+            "kl_channelwise": kl_channelwise,  # (batch, ch)
+        }  # TODO revisit, check dims
+        return kl_dict
+
+    def process_p_params(
+        self, p_params: torch.Tensor, var_clip_max: float
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.distributions.normal.Normal]:
+        """Process the input parameters to get the prior distribution p(z_i|z_{i+1}) (or p(z_L)).
+
+        Processing consists in:
+            - (optionally) 2D convolution on the input tensor to increase number of channels.
+            - split the resulting tensor into two chunks, the mean and the log-variance.
+            - (optionally) clip the log-variance to an upper threshold.
+            - define the normal distribution p(z) given the parameter tensors above.
+
+        Parameters
+        ----------
+        p_params: torch.Tensor
+            The input tensor to be processed.
+        var_clip_max: float
+            The maximum value reachable by the log-variance of the latent distribution.
+            Values exceeding this threshold are clipped.
+        """
+        if self.transform_p_params:
+            p_params = self.conv_in_p(p_params)
+        else:
+            assert p_params.size(1) == 2 * self.c_vars
+
+        # Define p(z)
+        p_mu, p_lv = p_params.chunk(2, dim=1)
+        if var_clip_max is not None:
+            p_lv = torch.clip(p_lv, max=var_clip_max)
+
+        p_mu = StableMean(p_mu)
+        p_lv = StableLogVar(p_lv, enable_stable=not self._use_naive_exponential)
+        p = Normal(p_mu.get(), p_lv.get_std())
+        return p_mu, p_lv, p
+
+    def process_q_params(
+        self, q_params: torch.Tensor, var_clip_max: float, allow_oddsizes: bool = False
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.distributions.normal.Normal]:
+        """
+        Process the input parameters to get the inference distribution q(z_i|z_{i+1}) (or q(z|x)).
+
+        Processing consists in:
+            - convolution on the input tensor to double the number of channels.
+            - split the resulting tensor into 2 chunks, respectively mean and log-var.
+            - (optionally) clip the log-variance to an upper threshold.
+            - (optionally) crop the resulting tensors to ensure that the last spatial dimension is even.
+            - define the normal distribution q(z) given the parameter tensors above.
+
+        Parameters
+        ----------
+        p_params: torch.Tensor
+            The input tensor to be processed.
+        var_clip_max: float
+            The maximum value reachable by the log-variance of the latent distribution.
+            Values exceeding this threshold are clipped.
+        """
+        q_params = self.conv_in_q(q_params)
+
+        q_mu, q_lv = q_params.chunk(2, dim=1)
+        if var_clip_max is not None:
+            q_lv = torch.clip(q_lv, max=var_clip_max)
+
+        if q_mu.shape[-1] % 2 == 1 and allow_oddsizes is False:
+            q_mu = F.center_crop(q_mu, q_mu.shape[-1] - 1)
+            q_lv = F.center_crop(q_lv, q_lv.shape[-1] - 1)
+            # TODO revisit ?!
+        q_mu = StableMean(q_mu)
+        q_lv = StableLogVar(q_lv, enable_stable=not self._use_naive_exponential)
+        q = Normal(q_mu.get(), q_lv.get_std())
+        return q_mu, q_lv, q
+
+    def forward(
+        self,
+        p_params: torch.Tensor,
+        q_params: Union[torch.Tensor, None] = None,
+        forced_latent: Union[torch.Tensor, None] = None,
+        force_constant_output: bool = False,
+        analytical_kl: bool = False,
+        mode_pred: bool = False,
+        use_uncond_mode: bool = False,
+        var_clip_max: Union[float, None] = None,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        """
+        Parameters
+        ----------
+        p_params: torch.Tensor
+            The output tensor of the top-down layer above (i.e., mu_{p,i+1}, sigma_{p,i+1}).
+        q_params: torch.Tensor, optional
+            The tensor resulting from merging the bu_value tensor at the same hierarchical level
+            from the bottom-up pass and the `p_params` tensor. Default is `None`.
+        forced_latent: torch.Tensor, optional
+            A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent
+            tensor and, hence, sampling does not happen. Default is `None`.
+        force_constant_output: bool, optional
+            Whether to copy the first sample (and rel. distrib parameters) over the whole batch.
+            This is used when doing experiment from the prior - q is not used.
+            Default is `False`.
+        analytical_kl: bool, optional
+            Whether to compute the KL divergence analytically or using Monte Carlo estimation.
+            Default is `False`.
+        mode_pred: bool, optional
+            Whether the model is in prediction mode. Default is `False`.
+        use_uncond_mode: bool, optional
+            Whether to use the uncoditional distribution p(z) to sample latents in prediction mode.
+            Default is `False`.
+        var_clip_max: float, optional
+            The maximum value reachable by the log-variance of the latent distribution.
+            Values exceeding this threshold are clipped. Default is `None`.
+        """
+        debug_qvar_max = 0
+
+        # Check sampling options consistency
+        assert forced_latent is None
+
+        # Get generative distribution p(z_i|z_{i+1})
+        p_mu, p_lv, p = self.process_p_params(p_params, var_clip_max)
+        p_params = (p_mu, p_lv)
+
+        if q_params is not None:
+            # Get inference distribution q(z_i|z_{i+1})
+            q_mu, q_lv, q = self.process_q_params(q_params, var_clip_max)
+            q_params = (q_mu, q_lv)
+            debug_qvar_max = torch.max(q_lv.get())
+            sampling_distrib = q
+            q_size = q_mu.get().shape[-1]
+            if p_mu.get().shape[-1] != q_size and mode_pred is False:
+                p_mu.centercrop_to_size(q_size)
+                p_lv.centercrop_to_size(q_size)
+        else:
+            sampling_distrib = p
+
+        # Sample latent variable
+        z = self.get_z(sampling_distrib, forced_latent, mode_pred, use_uncond_mode)
+
+        # TODO: not necessary, remove
+        # Copy one sample (and distrib parameters) over the whole batch.
+        # This is used when doing experiment from the prior - q is not used.
+        if force_constant_output:
+            z = z[0:1].expand_as(z).clone()
+            p_params = (
+                p_params[0][0:1].expand_as(p_params[0]).clone(),
+                p_params[1][0:1].expand_as(p_params[1]).clone(),
+            )
+
+        # Pass the sampled latent through the output convolution of stochastic block
+        out = self.conv_out(z)
+
+        if q_params is not None:
+            # Compute log q(z)
+            logprob_q = q.log_prob(z).sum(tuple(range(1, z.dim())))
+            # Compute KL divergence metrics
+            kl_dict = self.compute_kl_metrics(
+                p, p_params, q, q_params, mode_pred, analytical_kl, z
+            )
+        else:
+            kl_dict = {}
+            logprob_q = None
+
+        # Store meaningful quantities for later computation
+        data = kl_dict
+        data["z"] = z  # sampled variable at this layer (B, C, [Z], Y, X)
+        data["p_params"] = p_params  # (B, C, [Z], Y, X) where B is 1 or batch size
+        data["q_params"] = q_params  # (B, C, [Z], Y, X)
+        data["logprob_q"] = logprob_q  # (B, )
+        data["qvar_max"] = debug_qvar_max
+        return out, data

careamics/models/lvae/utils.py

@@ -2,7 +2,7 @@
 Script for utility functions needed by the LVAE model.
 """
 
-from typing import Iterable
+from typing import Literal, Sequence
 
 import numpy as np
 import torch
@@ -15,11 +15,6 @@ def torch_nanmean(inp):
     return torch.mean(inp[~inp.isnan()])
 
 
-def compute_batch_mean(x):
-    N = len(x)
-    return x.view(N, -1).mean(dim=1)
-
-
 def power_of_2(self, x):
     assert isinstance(x, int)
     if x == 1:
@@ -100,46 +95,76 @@ class ModelType(Enum):
     LadderVAETwoDataSetFinetuning = 28
 
 
-def _pad_crop_img(x, size, mode) -> torch.Tensor:
+def _pad_crop_img(
+    x: torch.Tensor, size: Sequence[int], mode: Literal["crop", "pad"]
+) -> torch.Tensor:
     """Pads or crops a tensor.
-    Pads or crops a tensor of shape (batch, channels, h, w) to new height
-    and width given by a tuple.
-    Args:
-        x (torch.Tensor): Input image
-        size (list or tuple): Desired size (height, width)
-        mode (str): Mode, either 'pad' or 'crop'
+
+    Pads or crops a tensor of shape (B, C, [Z], Y, X) to new shape.
+
+    Parameters:
+    -----------
+    x: torch.Tensor
+        Input image of shape (B, C, [Z], Y, X)
+    size: Sequence[int]
+        Desired size ([Z*], Y*, X*)
+    mode: Literal["crop", "pad"]
+        Mode, either 'pad' or 'crop'
+
     Returns:
+    --------
+    torch.Tensor:
         The padded or cropped tensor
     """
-    assert x.dim() == 4 and len(size) == 2
+    # TODO: Support cropping/padding on selected dimensions
+    assert (x.dim() == 4 and len(size) == 2) or (x.dim() == 5 and len(size) == 3)
+
     size = tuple(size)
-    x_size = x.size()[2:4]
+    x_size = x.size()[2:]
+
     if mode == "pad":
-        cond = x_size[0] > size[0] or x_size[1] > size[1]
+        cond = any(x_size[i] > size[i] for i in range(len(size)))
     elif mode == "crop":
-        cond = x_size[0] < size[0] or x_size[1] < size[1]
-    else:
-        raise ValueError(f"invalid mode '{mode}'")
+        cond = any(x_size[i] < size[i] for i in range(len(size)))
+
     if cond:
-        raise ValueError(f"trying to {mode} from size {x_size} to size {size}")
-    dr, dc = (abs(x_size[0] - size[0]), abs(x_size[1] - size[1]))
-    dr1, dr2 = dr // 2, dr - (dr // 2)
-    dc1, dc2 = dc // 2, dc - (dc // 2)
+        raise ValueError(f"Trying to {mode} from size {x_size} to size {size}")
+
+    diffs = [abs(x - s) for x, s in zip(x_size, size)]
+    d1 = [d // 2 for d in diffs]
+    d2 = [d - (d // 2) for d in diffs]
+
     if mode == "pad":
-        return nn.functional.pad(x, [dc1, dc2, dr1, dr2, 0, 0, 0, 0])
+        if x.dim() == 4:
+            padding = [d1[1], d2[1], d1[0], d2[0], 0, 0, 0, 0]
+        elif x.dim() == 5:
+            padding = [d1[2], d2[2], d1[1], d2[1], d1[0], d2[0], 0, 0, 0, 0]
+        return nn.functional.pad(x, padding)
     elif mode == "crop":
-        return x[:, :, dr1 : x_size[0] - dr2, dc1 : x_size[1] - dc2]
+        if x.dim() == 4:
+            return x[:, :, d1[0] : (x_size[0] - d2[0]), d1[1] : (x_size[1] - d2[1])]
+        elif x.dim() == 5:
+            return x[
+                :,
+                :,
+                d1[0] : (x_size[0] - d2[0]),
+                d1[1] : (x_size[1] - d2[1]),
+                d1[2] : (x_size[2] - d2[2]),
+            ]
 
 
-def pad_img_tensor(x, size) -> torch.Tensor:
-    """Pads a tensor.
-    Pads a tensor of shape (batch, channels, h, w) to a desired height and width.
-    Args:
-        x (torch.Tensor): Input image
-        size (list or tuple): Desired size (height, width)
+def pad_img_tensor(x: torch.Tensor, size: Sequence[int]) -> torch.Tensor:
+    """Pads a tensor
 
-    Returns
-    -------
+    Pads a tensor of shape (B, C, [Z], Y, X) to desired spatial dimensions.
+
+    Parameters:
+    -----------
+        x (torch.Tensor): Input image of shape (B, C, [Z], Y, X)
+        size (list or tuple): Desired size  ([Z*], Y*, X*)
+
+    Returns:
+    --------
         The padded tensor
     """
     return _pad_crop_img(x, size, "pad")
@@ -251,7 +276,15 @@ class StableLogVar:
     def get_std(self) -> torch.Tensor:
         return torch.sqrt(self.get_var())
 
-    def centercrop_to_size(self, size: Iterable[int]) -> None:
+    @property
+    def is_3D(self) -> bool:
+        """Check if the _lv tensor is 3D.
+
+        Recall that, in this framework, tensors have shape (B, C, [Z], Y, X).
+        """
+        return self._lv.dim() == 5
+
+    def centercrop_to_size(self, size: Sequence[int]) -> None:
         """
         Centercrop the log-variance tensor to the desired size.
 
@@ -260,6 +293,8 @@ class StableLogVar:
         size: torch.Tensor
             The desired size of the log-variance tensor.
         """
+        assert not self.is_3D, "Centercrop is implemented only for 2D tensors."
+
         if self._lv.shape[-1] == size:
             return
 
@@ -276,15 +311,26 @@ class StableMean:
     def get(self) -> torch.Tensor:
         return self._mean
 
-    def centercrop_to_size(self, size: Iterable[int]) -> None:
+    @property
+    def is_3D(self) -> bool:
+        """Check if the _mean tensor is 3D.
+
+        Recall that, in this framework, tensors have shape (B, C, [Z], Y, X).
         """
-        Centercrop the mean tensor to the desired size.
+        return self._mean.dim() == 5
+
+    def centercrop_to_size(self, size: Sequence[int]) -> None:
+        """Centercrop the mean tensor to the desired size.
+
+        Implemented only in the case of 2D tensors.
 
         Parameters
         ----------
         size: torch.Tensor
             The desired size of the log-variance tensor.
         """
+        assert not self.is_3D, "Centercrop is implemented only for 2D tensors."
+
         if self._mean.shape[-1] == size:
             return
 
@@ -356,40 +402,3 @@ def kl_normal_mc(z, p_mulv, q_mulv):
     p_distrib = Normal(p_mu.get(), p_std)
     q_distrib = Normal(q_mu.get(), q_std)
     return q_distrib.log_prob(z) - p_distrib.log_prob(z)
-
-
-def free_bits_kl(
-    kl: torch.Tensor, free_bits: float, batch_average: bool = False, eps: float = 1e-6
-) -> torch.Tensor:
-    """
-    Computes free-bits version of KL divergence.
-    Ensures that the KL doesn't go to zero for any latent dimension.
-    Hence, it contributes to use latent variables more efficiently,
-    leading to better representation learning.
-
-    NOTE:
-    Takes in the KL with shape (batch size, layers), returns the KL with
-    free bits (for optimization) with shape (layers,), which is the average
-    free-bits KL per layer in the current batch.
-    If batch_average is False (default), the free bits are per layer and
-    per batch element. Otherwise, the free bits are still per layer, but
-    are assigned on average to the whole batch. In both cases, the batch
-    average is returned, so it's simply a matter of doing mean(clamp(KL))
-    or clamp(mean(KL)).
-
-    Args:
-        kl (torch.Tensor)
-        free_bits (float)
-        batch_average (bool, optional))
-        eps (float, optional)
-
-    Returns
-    -------
-        The KL with free bits
-    """
-    assert kl.dim() == 2
-    if free_bits < eps:
-        return kl.mean(0)
-    if batch_average:
-        return kl.mean(0).clamp(min=free_bits)
-    return kl.clamp(min=free_bits).mean(0)

careamics/utils/lightning_utils.py

@@ -0,0 +1,57 @@
+"""PyTorch lightning utilities."""
+
+from pathlib import Path
+from typing import Union
+
+
+def read_csv_logger(experiment_name: str, log_folder: Union[str, Path]) -> dict:
+    """Return the loss curves from the csv logs.
+
+    Parameters
+    ----------
+    experiment_name : str
+        Name of the experiment.
+    log_folder : Path or str
+        Path to the folder containing the csv logs.
+
+    Returns
+    -------
+    dict
+        Dictionary containing the loss curves, with keys "train_epoch", "val_epoch",
+        "train_loss" and "val_loss".
+    """
+    path = Path(log_folder) / experiment_name
+
+    # find the most recent of version_* folders
+    versions = [int(v.name.split("_")[-1]) for v in path.iterdir() if v.is_dir()]
+    version = max(versions)
+
+    path_log = path / f"version_{version}" / "metrics.csv"
+
+    epochs = []
+    train_losses_tmp = []
+    val_losses_tmp = []
+    with open(path_log) as f:
+        lines = f.readlines()
+
+        for single_line in lines[1:]:
+            epoch, _, train_loss, _, val_loss = single_line.strip().split(",")
+
+            epochs.append(epoch)
+            train_losses_tmp.append(train_loss)
+            val_losses_tmp.append(val_loss)
+
+    # train and val are not logged on the same row and can have different lengths
+    train_epoch = [
+        int(epochs[i]) for i in range(len(epochs)) if train_losses_tmp[i] != ""
+    ]
+    val_epoch = [int(epochs[i]) for i in range(len(epochs)) if val_losses_tmp[i] != ""]
+    train_losses = [float(loss) for loss in train_losses_tmp if loss != ""]
+    val_losses = [float(loss) for loss in val_losses_tmp if loss != ""]
+
+    return {
+        "train_epoch": train_epoch,
+        "val_epoch": val_epoch,
+        "train_loss": train_losses,
+        "val_loss": val_losses,
+    }

careamics/utils/serializers.py

@@ -21,6 +21,8 @@ def _array_to_json(arr: Union[np.ndarray, torch.Tensor]) -> str:
     str
         JSON string representing the array.
     """
+    if isinstance(arr, str):
+        return arr
     return json.dumps(arr.tolist())
 
 

careamics/utils/torch_utils.py

@@ -84,7 +84,7 @@ def get_optimizers() -> Dict[str, str]:
 def get_scheduler(
     name: str,
 ) -> Union[
-    torch.optim.lr_scheduler.LRScheduler,
+    # torch.optim.lr_scheduler.LRScheduler,
     torch.optim.lr_scheduler.ReduceLROnPlateau,
 ]:
     """