modify logic to call prepare and setup during instantiation of data module

- Lightning-AI/pytorch-lightning#20602 (comment)

Author: aditya0by0

chebai/cli.py

@@ -38,15 +38,26 @@ def add_arguments_to_parser(self, parser: LightningArgumentParser):
         Args:
             parser (LightningArgumentParser): Argument parser instance.
         """
-        # for kind in ("train", "val", "test"):
-        #     for average in ("micro-f1", "macro-f1", "balanced-accuracy"):
-        #         parser.link_arguments(
-        #             "model.init_args.out_dim",
-        #             f"model.init_args.{kind}_metrics.init_args.metrics.{average}.init_args.num_labels",
-        #         )
-        # parser.link_arguments(
-        #     "model.init_args.out_dim", "trainer.callbacks.init_args.num_labels"
-        # )
+
+        parser.link_arguments(
+            "data.num_of_labels", "model.init_args.out_dim", apply_on="instantiate"
+        )
+        parser.link_arguments(
+            "data.feature_vector_size",
+            "model.init_args.input_dim",
+            apply_on="instantiate",
+        )
+
+        for kind in ("train", "val", "test"):
+            for average in ("micro-f1", "macro-f1", "balanced-accuracy"):
+                parser.link_arguments(
+                    "data.num_of_labels",
+                    f"model.init_args.{kind}_metrics.init_args.metrics.{average}.init_args.num_labels",
+                    apply_on="instantiate",
+                )
+        parser.link_arguments(
+            "data.num_of_labels", "trainer.callbacks.init_args.num_labels"
+        )
         parser.link_arguments(
             "data", "model.init_args.criterion.init_args.data_extractor"
         )

chebai/models/base.py

@@ -35,6 +35,8 @@ class ChebaiBaseNet(LightningModule):
     def __init__(
         self,
         criterion: torch.nn.Module = None,
+        out_dim: Optional[int] = None,
+        input_dim: Optional[int] = None,
         train_metrics: Optional[torch.nn.Module] = None,
         val_metrics: Optional[torch.nn.Module] = None,
         test_metrics: Optional[torch.nn.Module] = None,
@@ -47,7 +49,12 @@ def __init__(
         self.save_hyperparameters(
             ignore=["criterion", "train_metrics", "val_metrics", "test_metrics"]
         )
-        self.out_dim = None
+
+        self.out_dim = out_dim
+        self.input_dim = input_dim
+        assert out_dim is not None, "out_dim must be specified"
+        assert input_dim is not None, "input_dim must be specified"
+
         if optimizer_kwargs:
             self.optimizer_kwargs = optimizer_kwargs
         else:
@@ -69,14 +76,6 @@ def __init_subclass__(cls, **kwargs):
         else:
             _MODEL_REGISTRY[cls.NAME] = cls
 
-    def setup(self, stage: str) -> None:
-        if self.trainer and hasattr(self.trainer, "datamodule"):
-            self.out_dim = int(self.trainer.datamodule.hparams.num_of_labels)
-        else:
-            raise ValueError("Trainer has no data module")
-        assert self.out_dim is not None, "Model output dimension is None"
-        print(f"Output Dimension for the model: {self.out_dim}")
-
     def _get_prediction_and_labels(
         self, data: Dict[str, Any], labels: torch.Tensor, output: torch.Tensor
     ) -> (torch.Tensor, torch.Tensor):

chebai/models/ffn.py

@@ -20,15 +20,9 @@ def __init__(
         **kwargs
     ):
         super().__init__(**kwargs)
-        self.input_size = input_size
-        self.hidden_layers = hidden_layers
-
-    def setup(self, stage: str) -> None:
-        super().setup(stage)
-
         layers = []
-        current_layer_input_size = self.input_size
-        for hidden_dim in self.hidden_layers:
+        current_layer_input_size = input_size
+        for hidden_dim in hidden_layers:
             layers.append(MLPBlock(current_layer_input_size, hidden_dim))
             layers.append(Residual(MLPBlock(hidden_dim, hidden_dim)))
             current_layer_input_size = hidden_dim

chebai/preprocessing/datasets/base.py

@@ -117,6 +117,25 @@ def __init__(
             os.makedirs(os.path.join(self.processed_dir, self.fold_dir), exist_ok=True)
         self.save_hyperparameters()
 
+        self._num_of_labels = None
+        self._feature_vector_size = None
+        self._prepare_data_flag = 1
+        self._setup_data_flag = 1
+        self.prepare_data()
+        self.setup()
+
+    @property
+    def num_of_labels(self):
+        assert self._num_of_labels is not None, "num of labels must be set"
+        return self._num_of_labels
+
+    @property
+    def feature_vector_size(self):
+        assert (
+            self._feature_vector_size is not None
+        ), "size of feature vector must be set"
+        return self._feature_vector_size
+
     @property
     def identifier(self) -> tuple:
         """Identifier for the dataset."""
@@ -381,6 +400,12 @@ def predict_dataloader(
         """
         return self.dataloader(self.prediction_kind, shuffle=False, **kwargs)
 
+    def prepare_data(self) -> None:
+        if self._prepare_data_flag != 1:
+            return
+
+        self._prepare_data_flag += 1
+
     def setup(self, **kwargs):
         """
         Setup the data module.
@@ -390,6 +415,11 @@ def setup(self, **kwargs):
         Args:
             **kwargs: Additional keyword arguments.
         """
+        if self._setup_data_flag != 1:
+            return
+
+        self._setup_data_flag += 1
+
         rank_zero_info(f"Check for processed data in {self.processed_dir}")
         rank_zero_info(f"Cross-validation enabled: {self.use_inner_cross_validation}")
         if any(
@@ -401,20 +431,20 @@ def setup(self, **kwargs):
         if not ("keep_reader" in kwargs and kwargs["keep_reader"]):
             self.reader.on_finish()
 
-        self._add_num_of_labels_to_hparams()
+        self._set_processed_data_props()
 
-    def _add_num_of_labels_to_hparams(self):
-        num_of_labels = len(
-            torch.load(
-                os.path.join(
-                    self.processed_dir, self.processed_file_names_dict["data"]
-                ),
-                weights_only=False,
-            )[0]["labels"]
-        )
+    def _set_processed_data_props(self):
 
-        print(f"Number of labels for loaded data: {num_of_labels}")
-        self.hparams.num_of_labels = num_of_labels
+        single_data_instance = torch.load(
+            os.path.join(self.processed_dir, self.processed_file_names_dict["data"]),
+            weights_only=False,
+        )[0]
+
+        self._num_of_labels = len(single_data_instance["labels"])
+        self._feature_vector_size = len(single_data_instance["features"])
+
+        print(f"Number of labels for loaded data: {self._num_of_labels}")
+        print(f"Feature vector size: {self._feature_vector_size}")
 
     def setup_processed(self):
         """
@@ -541,6 +571,7 @@ def prepare_data(self):
         """
         Placeholder for data preparation logic.
         """
+        super().prepare_data()
         for s in self.subsets:
             s.prepare_data()
 
@@ -553,10 +584,14 @@ def setup(self, **kwargs):
         Args:
             **kwargs: Additional keyword arguments.
         """
+        if self._setup_data_flag != 1:
+            return
+
+        self._setup_data_flag += 1
         for s in self.subsets:
             s.setup(**kwargs)
 
-        self._add_num_of_labels_to_hparams()
+        self._set_processed_data_props()
 
     def dataloader(self, kind: str, **kwargs) -> DataLoader:
         """
@@ -752,6 +787,7 @@ def prepare_data(self, *args: Any, **kwargs: Any) -> None:
         Returns:
             None
         """
+        super().prepare_data()
         print("Checking for processed data in", self.processed_dir_main)
 
         processed_name = self.processed_main_file_names_dict["data"]

chebai/preprocessing/datasets/chebi.py

@@ -60,6 +60,7 @@ def raw_file_names(self):
         return ["test.pkl", "train.pkl", "validation.pkl"]
 
     def prepare_data(self, *args, **kwargs):
+        super().prepare_data()
         print("Check for raw data in", self.raw_dir)
         if any(
             not os.path.isfile(os.path.join(self.raw_dir, f))

chebai/preprocessing/datasets/deepGO/go_uniprot.py

@@ -783,6 +783,11 @@ def prepare_data(self, *args: Any, **kwargs: Any) -> None:
         Raises:
             FileNotFoundError: If the processed data file does not exist.
         """
+        if self._prepare_data_flag != 1:
+            return
+
+        self._prepare_data_flag += 1
+
         print("Checking for processed data in", self.processed_dir_main)
 
         processed_name = self.processed_main_file_names_dict["data"]

chebai/preprocessing/datasets/deepGO/protein_pretraining.py

@@ -64,6 +64,11 @@ def prepare_data(self, *args: Any, **kwargs: Any) -> None:
             *args: Additional positional arguments.
             **kwargs: Additional keyword arguments.
         """
+        if self._prepare_data_flag != 1:
+            return
+
+        self._prepare_data_flag += 1
+
         processed_name = self.processed_main_file_names_dict["data"]
         if not os.path.isfile(os.path.join(self.processed_dir_main, processed_name)):
             print("Missing processed data file (`data.pkl` file)")

chebai/preprocessing/datasets/pubchem.py

@@ -183,6 +183,7 @@ def prepare_data(self, *args, **kwargs):
         """
         Checks for raw data and downloads if necessary.
         """
+        super().prepare_data()
         print("Check for raw data in", self.raw_dir)
         if any(
             not os.path.isfile(os.path.join(self.raw_dir, f))

chebai/preprocessing/datasets/tox21.py

@@ -118,6 +118,10 @@ def setup_processed(self) -> None:
 
     def setup(self, **kwargs) -> None:
         """Sets up the dataset by downloading and processing if necessary."""
+        if self._setup_data_flag != 1:
+            return
+
+        self._setup_data_flag += 1
         if any(
             not os.path.isfile(os.path.join(self.raw_dir, f))
             for f in self.raw_file_names
@@ -129,7 +133,7 @@ def setup(self, **kwargs) -> None:
         ):
             self.setup_processed()
 
-        self._add_num_of_labels_to_hparams()
+        self._set_processed_data_props()
 
     def _load_data_from_file(self, input_file_path: str) -> List[Dict]:
         """Loads data from a CSV file.
@@ -302,6 +306,10 @@ def setup_processed(self) -> None:
 
     def setup(self, **kwargs) -> None:
         """Sets up the dataset by downloading and processing if necessary."""
+        if self._setup_data_flag != 1:
+            return
+
+        self._setup_data_flag += 1
         if any(
             not os.path.isfile(os.path.join(self.raw_dir, f))
             for f in self.raw_file_names
@@ -313,7 +321,7 @@ def setup(self, **kwargs) -> None:
         ):
             self.setup_processed()
 
-        self._add_num_of_labels_to_hparams()
+        self._set_processed_data_props()
 
     def _load_dict(self, input_file_path: str) -> Generator[Dict, None, None]:
         """Loads data from a CSV file as a generator.