classifier

Author: Dingel321

pyproject.toml

@@ -28,4 +28,5 @@ dependencies = [
 
 [project.scripts]
 train_npe_model = "cryo_sbi.inference.command_line_tools:cl_npe_train_no_saving"
+train_classifier = "cryo_sbi.inference.command_line_tools:cl_classifier_train_no_saving"
 model_to_tensor = "cryo_sbi.utils.command_line_tools:cl_models_to_tensor"

src/cryo_sbi/inference/command_line_tools.py

@@ -1,10 +1,10 @@
 import argparse
-from cryo_sbi.inference.train_npe_model import (
-    npe_train_no_saving,
+from cryo_sbi.inference.train_models import (
+    train_classifier
 )
 
 
-def cl_npe_train_no_saving():
+def cl_classifier_train_no_saving():
     cl_parser = argparse.ArgumentParser()
 
     cl_parser.add_argument(
@@ -47,7 +47,7 @@ def cl_npe_train_no_saving():
 
     args = cl_parser.parse_args()
 
-    npe_train_no_saving(
+    train_classifier(
         image_config=args.image_config_file,
         train_config=args.train_config_file,
         epochs=args.epochs,
@@ -59,4 +59,4 @@ def cl_npe_train_no_saving():
         device=args.train_device,
         saving_frequency=args.saving_freq,
         simulation_batch_size=args.simulation_batch_size,
-    )
+    )

src/cryo_sbi/inference/models/build_models.py

@@ -6,7 +6,7 @@
 from cryo_sbi.inference.models.embedding_nets import EMBEDDING_NETS
 
 
-def build_npe_flow_model(config: dict, **embedding_kwargs) -> nn.Module:
+def build_classifier(config: dict, **embedding_kwargs) -> nn.Module:
     """
     Function to build NPE estimator with embedding net
     from config_file
@@ -19,17 +19,6 @@ def build_npe_flow_model(config: dict, **embedding_kwargs) -> nn.Module:
         estimator (nn.Module): NPE estimator
     """
 
-    if config["MODEL"] == "MAF":
-        model = zuko.flows.MAF
-    elif config["MODEL"] == "NSF":
-        model = zuko.flows.NSF
-    elif config["MODEL"] == "SOSPF":
-        model = zuko.flows.SOSPF
-    else:
-        raise NotImplementedError(
-            f"Model : {config['MODEL']} has not been implemented yet!"
-        )
-
     try:
         embedding = partial(
             EMBEDDING_NETS[config["EMBEDDING"]], config["OUT_DIM"], **embedding_kwargs
@@ -40,20 +29,16 @@ def build_npe_flow_model(config: dict, **embedding_kwargs) -> nn.Module:
 The following embeddings are implemented : {[key for key in EMBEDDING_NETS.keys()]}"
         )
 
-    estimator = estimator_models.NPEWithEmbedding(
+    estimator = estimator_models.ClassifierWithEmbedding(
         embedding_net=embedding,
         output_embedding_dim=config["OUT_DIM"],
-        num_transforms=config["NUM_TRANSFORM"],
-        num_hidden_flow=config["NUM_HIDDEN_FLOW"],
-        hidden_flow_dim=config["HIDDEN_DIM_FLOW"],
-        flow=model,
-        theta_shift=config["THETA_SHIFT"],
-        theta_scale=config["THETA_SCALE"],
-        **{"activation": partial(nn.LeakyReLU, 0.1)},
+        num_classes=config["NUM_CLASSES"],
+        num_layers=config["NUM_LAYERS"],
+        nodes_per_layer=config["NODES_PER_LAYER"],
+        **{
+            "activation": partial(nn.LeakyReLU, 0.1),
+            "dropout": config["DROPOUT"],
+        },
     )
 
     return estimator
-
-
-def build_nre_classifier_model(config: dict, **embedding_kwargs) -> nn.Module:
-    raise NotImplementedError("NRE classifier model has not been implemented yet!")

src/cryo_sbi/inference/models/estimator_models.py

@@ -1,76 +1,46 @@
+from typing import Tuple
 import torch
 import torch.nn as nn
 import zuko
 from lampe.inference import NPE, NRE
 
 
-class Standardize(nn.Module):
-    """
-    Module to standardize inputs and retransform them to the original space
-
-    Args:
-        mean (torch.Tensor): mean of the data
-        std (torch.Tensor): standard deviation of the data
-
-    Returns:
-        standardized (torch.Tensor): standardized data
-    """
-
-    # Code adapted from :https://github.com/mackelab/sbi/blob/main/sbi/utils/sbiutils.py
-    def __init__(self, mean: float, std: float) -> None:
-        super(Standardize, self).__init__()
-        mean, std = map(torch.as_tensor, (mean, std))
-        self.mean = mean
-        self.std = std
-        self.register_buffer("_mean", mean)
-        self.register_buffer("_std", std)
-
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        """
-        Standardize the input tensor
-
-        Args:
-            tensor (torch.Tensor): input tensor
-
-        Returns:
-            standardized (torch.Tensor): standardized tensor
-        """
-
-        return (tensor - self._mean) / self._std
-
-    def transform(self, tensor: torch.Tensor) -> torch.Tensor:
-        """
-        Transform the standardized tensor back to the original space
-
-        Args:
-            tensor (torch.Tensor): input tensor
+class Classifier(nn.Module):
+    def __init__(self, input_dim, out_dim, num_layers, nodes_per_layer, activation=nn.ReLU, dropout=0.0):
+        super().__init__()
+        self.classifier = nn.ModuleList()
 
-        Returns:
-            retransformed (torch.Tensor): retransformed tensor
-        """
+        for i in range(num_layers):
+            if i == 0:
+                self.classifier.append(nn.Linear(input_dim, nodes_per_layer))
+            else:
+                self.classifier.append(nn.Linear(nodes_per_layer, nodes_per_layer))
+                if dropout > 0.0:
+                    self.classifier.append(nn.Dropout(dropout))
+            self.classifier.append(activation())
+        
+        self.classifier.append(nn.Linear(nodes_per_layer, out_dim))
+        self.classifier = nn.Sequential(*self.classifier)
 
-        return (tensor * self._std) + self._mean
+    def forward(self, x):
+        return self.classifier(x)
 
 
-class NPEWithEmbedding(nn.Module):
-    """Neural Posterior Estimation with embedding net
+class ClassifierWithEmbedding(nn.Module):
+    """Classification with embedding net
 
     Attributes:
-        npe (NPE): NPE model
+        classifier: Classification model
         embedding (nn.Module): embedding net
-        standardize (Standardize): standardization module
     """
 
     def __init__(
         self,
         embedding_net: nn.Module,
         output_embedding_dim: int,
-        num_transforms: int = 4,
-        num_hidden_flow: int = 2,
-        hidden_flow_dim: int = 128,
-        flow: nn.Module = zuko.flows.MAF,
-        theta_shift: float = 0.0,
-        theta_scale: float = 1.0,
+        num_classes: int = 2,
+        num_layers: int = 5,
+        nodes_per_layer: int = 128,
         **kwargs,
     ) -> None:
         """
@@ -93,55 +63,50 @@ def __init__(
 
         super().__init__()
 
-        self.npe = NPE(
-            1,
+        self.classifier = Classifier(
             output_embedding_dim,
-            transforms=num_transforms,
-            build=flow,
-            hidden_features=[*[hidden_flow_dim] * num_hidden_flow, 128, 64],
+            num_classes,
+            num_layers,
+            nodes_per_layer,
             **kwargs,
         )
-
         self.embedding = embedding_net()
-        self.standardize = Standardize(theta_shift, theta_scale)
 
-    def forward(self, theta: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
-        Forward pass of the NPE model
-
+        Forward pass of the classifier model
         Args:
             theta (torch.Tensor): Conformational parameters.
             x (torch.Tensor): Image to condition the posterior on.
-
         Returns:
-            torch.Tensor: Log probability of the posterior.
+            torch.Tensor: unnormalized class probabilities.
         """
 
-        return self.npe(self.standardize(theta), self.embedding(x))
-
-    def flow(self, x: torch.Tensor):
+        return self.classifier(self.embedding(x))
+    
+    def prob(self, x: torch.Tensor) -> torch.Tensor:
         """
-        Conditions the posterior on an image.
-
+        Predict the class probabilities for the input data.
+        
         Args:
-            x (torch.Tensor): Image to condition the posterior on.
-
+            x (torch.Tensor): Input data.
+        
         Returns:
-            zuko.flows.Flow: The posterior distribution.
+            torch.Tensor: Class probabilities.
         """
-        return self.npe.flow(self.embedding(x))
-
-    def sample(self, x: torch.Tensor, shape=(1,)) -> torch.Tensor:
+        return torch.nn.functional.softmax(self.forward(x))
+    
+    def logits_embedding(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         """
-        Generate samples from the posterior distribution.
+        Get the logits from the classifier and the embedding from the embedding net.
 
         Args:
-            x (torch.Tensor): Image to condition the posterior on.
-            shape (tuple, optional): Shape of the samples. Defaults to (1,).
+            x (torch.Tensor): Input data.
 
         Returns:
-            torch.Tensor: Samples from the posterior distribution.
+            torch.Tensor: Logits from the classifier.
+            torch.Tensor: Embedding from the embedding net.
         """
-
-        samples_standardized = self.flow(x).sample(shape)
-        return self.standardize.transform(samples_standardized)
+        embeddings = self.embedding(x)
+        logits = self.classifier(embeddings)
+        return logits, embeddings

src/cryo_sbi/inference/priors.py

@@ -23,6 +23,31 @@ def gen_quat() -> torch.Tensor:
     return quat
 
 
+class IndexPrior():
+    def __init__(self, max_index: int, device="cpu") -> None:
+        self.max_index = max_index
+        self.device = device
+
+        self.index_prior = torch.distributions.Categorical(
+            probs=torch.tensor(
+                [1 / (max_index + 1) for _ in range(max_index + 1)],
+                device=device,
+            )
+        )
+
+    def sample(self, shape) -> torch.Tensor:
+        """
+        Sample indices from the prior distribution.
+
+        Args:
+            shape (tuple): Shape of the samples to be generated.
+
+        Returns:
+            torch.Tensor: Sampled indices.
+        """
+        return self.index_prior.sample(shape)
+
+
 def get_image_priors(
     max_index, image_config: dict, device="cuda"
 ) -> zuko.distributions.BoxUniform:
@@ -110,10 +135,7 @@ def get_image_priors(
         ndims=1,
     )
 
-    index_prior = zuko.distributions.BoxUniform(
-        lower=torch.tensor([0], dtype=torch.float32, device=device),
-        upper=torch.tensor([max_index], dtype=torch.float32, device=device),
-    )
+    index_prior = IndexPrior(max_index, device)
     quaternion_prior = QuaternionPrior(device)
     if (
         image_config.get("ROTATIONS")
@@ -132,7 +154,6 @@ def get_image_priors(
         b_factor_prior,
         amp_prior,
         snr_prior,
-        device=device,
     )
 
 
@@ -168,7 +189,6 @@ def __init__(
         b_factor_prior,
         amp_prior,
         snr_prior,
-        device,
     ) -> None:
         self.priors = [
             index_prior,