Merge branch 'training'

Author: lus105

configs/data/classification_dir.yaml

@@ -26,5 +26,4 @@ val_test_transforms:
       size: ${data.image_size}
     - _target_: torchvision.transforms.ToTensor
 
-save_predict_images: false
-num_classes: ???
+save_predict_images: false

configs/experiment/train_cnn.yaml

@@ -20,17 +20,14 @@ trainer:
 data:
   batch_size: 32
   num_workers: 1
-  num_classes: 6
 
 model:
   net:
     _target_: src.models.components.base_model.BaseModel
     model_name: timm/mobilenetv3_large_100.ra_in1k
     pretrained: True
-    num_classes: ${data.num_classes}
   loss:
     _target_: torch.nn.CrossEntropyLoss
   ckpt_path: null
-  num_classes: ${data.num_classes}
 
 export_to_onnx: True

configs/experiment/train_cnn_multi.yaml

@@ -20,10 +20,8 @@ trainer:
 data:
   batch_size: 8
   num_workers: 1
-  num_classes: 2
 
 model:
   net:
     multi_head: false
   ckpt_path: null
-  num_classes: ${data.num_classes}

configs/experiment/train_grain.yaml

@@ -12,11 +12,10 @@ tags: ["classification", "train", "cnn"]
 seed: 12345
 
 trainer:
-  min_epochs: 10
-  max_epochs: 20
-  gradient_clip_val: 0.5
+  min_epochs: 20
+  max_epochs: 35
   accelerator: gpu
-  precision: "bf16"
+  precision: "bf16-mixed"
 
 data:
   batch_size: 32
@@ -42,19 +41,15 @@ data:
         replicate_borders: false
         inpaint_color: [215, 215, 215]
       - _target_: torchvision.transforms.ToTensor
-      
-  num_classes: 6
 
 model:
   net:
     _target_: src.models.components.base_model.BaseModel
     model_name: timm/mobilenetv3_large_100.ra_in1k
     pretrained: True
-    num_classes: ${data.num_classes}
   loss:
     _target_: torch.nn.CrossEntropyLoss
   ckpt_path: null
-  num_classes: ${data.num_classes}
 
 callbacks:
   model_checkpoint:

configs/experiment/train_vit.yaml

@@ -19,18 +19,15 @@ trainer:
 data:
   batch_size: 64
   num_workers: 1
-  num_classes: 6
 
 model:
   net:
     _target_: src.models.components.base_model.BaseModel
     model_name: timm/vit_tiny_patch16_224.augreg_in21k_ft_in1k
     pretrained: True
-    num_classes: ${data.num_classes}
     img_size: ${data.image_size}
   loss:
     _target_: torch.nn.CrossEntropyLoss
   ckpt_path: null
-  num_classes: ${data.num_classes}
 
 export_to_onnx: True

configs/experiment/train_vit_multi.yaml

@@ -19,7 +19,6 @@ trainer:
 data:
   batch_size: 64
   num_workers: 1
-  num_classes: 2
 
 model:
   net:
@@ -28,4 +27,3 @@ model:
     lr: 5e-5
   compile: false
   ckpt_path: null
-  num_classes: ${data.num_classes}

configs/model/cnn.yaml

@@ -22,5 +22,3 @@ loss:
 compile: false
 
 ckpt_path: ???
-
-num_classes: ???

configs/model/vit.yaml

@@ -16,6 +16,4 @@ loss:
 
 compile: false
 
-ckpt_path: ???
-
-num_classes: ???
+ckpt_path: ???

src/data/classification_datamodule.py

@@ -25,7 +25,6 @@ def __init__(
         train_transforms: Compose = None,
         val_test_transforms: Compose = None,
         save_predict_images: bool = False,
-        num_classes: int = 2,
     ) -> None:
         """Initialize a `DirDataModule`.
 
@@ -41,7 +40,6 @@ def __init__(
             train_transforms (Compose, optional): Train split transformations. Defaults to None.
             val_test_transforms (Compose, optional): Validation and test split transformations. Defaults to None.
             save_predict_images (bool, optional): Save images in predict mode? Defaults to False.
-            num_classes (int, optional): Number of classes in the dataset.
         """
         super().__init__()
 
@@ -55,29 +53,31 @@ def __init__(
         self.train_transforms = train_transforms
         self.val_test_transforms = val_test_transforms
         self.save_predict_images = save_predict_images
-        self._num_classes = num_classes
         self.channels = channels
         self._class_names: Optional[list[str]] = None
         self.data_train: Optional[Dataset] = None
         self.data_val: Optional[Dataset] = None
         self.data_test: Optional[Dataset] = None
         self.data_predict: Optional[Dataset] = None
+        self.setup_stages_done = set()
 
     @property
     def num_classes(self) -> int:
         """Get the number of classes.
 
         Returns:
-            int: The number of classes (2).
+            int: The number of classes.
         """
-        return self._num_classes
+        if self._class_names is None and self.data_train is None:
+            self.setup(stage='fit')
+
+        return len(self._class_names)
 
     @property
-    def class_names(self):
+    def class_names(self) -> Optional[list[str]]:
         """Automatically extract class names from the dataset."""
-
-        if self._class_names is None and hasattr(self.data_train, 'classes'):
-            self._class_names = self.data_train.classes
+        if self._class_names is None and self.data_train is None:
+            self.setup(stage='fit')
 
         return self._class_names
 
@@ -102,8 +102,10 @@ def setup(self, stage: Optional[str] = None) -> None:
             stage (Optional[str], optional): The stage to setup. Either `"fit"`,
             `"validate"`, `"test"`, or `"predict"`. Defaults to None.
         """
-
         if stage in {'fit', 'validate', 'test'}:
+            if 'fit' in self.setup_stages_done:
+                return
+
             self.data_train = ImageFolder(
                 root=Path(self.train_data_dir),
                 transform=self.train_transforms,
@@ -118,12 +120,26 @@ def setup(self, stage: Optional[str] = None) -> None:
                 root=Path(self.val_data_dir),
                 transform=self.val_test_transforms,
             )
+
+            if hasattr(self.data_train, 'classes') and self._class_names is None:
+                self._class_names = self.data_train.classes
+            
+            self.setup_stages_done.add('fit')
+
         elif stage == 'predict':
+            if 'predict' in self.setup_stages_done:
+                return
+            
             self.data_predict = ImageFolder(
                 root=Path(self.test_data_dir),
                 transform=self.val_test_transforms,
             )
 
+            if hasattr(self.data_predict, 'classes') and self._class_names is None:
+                self._class_names = self.data_predict.classes
+            
+            self.setup_stages_done.add('predict')
+
     def train_dataloader(self) -> DataLoader[Any]:
         """Create and return the train dataloader.
 

src/models/classification_module.py

@@ -17,7 +17,7 @@ def __init__(
         loss: torch.nn.modules.loss,
         compile: bool,
         ckpt_path: str,
-        num_classes: int = None,
+        num_classes: int,
     ) -> None:
         """Initialize lightning module.
 
@@ -28,7 +28,7 @@ def __init__(
             loss (torch.nn.modules.loss): Loss function.
             compile (bool): Compile model.
             ckpt_path (string): Model chekpoint path.
-            num_classes (int, optional): Number of classes.
+            num_classes (int): Number of classes.
         """
         super().__init__()
         # model
@@ -220,6 +220,7 @@ def setup(self, stage: str) -> None:
         Args:
             stage (str): Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
         """
+        self.net.load_model(num_classes=self.num_classes)
         if self.compile and stage == 'fit':
             self.net = torch.compile(self.net)
         if self.ckpt_path: