Fix issue in cifar10, add note about protocol

Signed-off-by: Fabrice Normandin <[email protected]>

project/datamodules/image_classification/cifar10.py

@@ -7,7 +7,6 @@
 from project.datamodules.image_classification.image_classification import (
     ImageClassificationDataModule,
 )
-from project.datamodules.vision import VisionDataModule
 from project.utils.typing_utils import C, H, W
 
 
@@ -40,7 +39,7 @@ def cifar10_unnormalization(x: torch.Tensor) -> torch.Tensor:
     return (x * std) + mean
 
 
-class CIFAR10DataModule(ImageClassificationDataModule, VisionDataModule):
+class CIFAR10DataModule(ImageClassificationDataModule):
     """
     .. figure:: https://3qeqpr26caki16dnhd19sv6by6v-wpengine.netdna-ssl.com/wp-content/uploads/2019/01/
         Plot-of-a-Subset-of-Images-from-the-CIFAR-10-Dataset.png

project/datamodules/image_classification/image_classification.py

@@ -12,11 +12,18 @@
 ImageBatchType = TypeVar("ImageBatchType", bound=tuple[Image, Tensor])
 
 
+# todo: this should probably be a protocol. The only issue with that is that we do `issubclass` in
+# tests to determine which datamodule configs are for image classification, so we can't do that
+# with a Protocol.
+
+
 class ImageClassificationDataModule(
     VisionDataModule[ImageBatchType], ClassificationDataModule[ImageBatchType]
 ):
     """Lightning data modules for image classification."""
 
+    # This just adds the `num_classes` property to `VisionDataModule`.
+
     num_classes: int
     """Number of classes in the dataset."""
 

project/datamodules/image_classification/imagenet.py

@@ -21,7 +21,9 @@
 from torchvision.models.resnet import ResNet152_Weights
 from torchvision.transforms import v2 as transforms
 
-from project.datamodules.vision import VisionDataModule
+from project.datamodules.image_classification.image_classification import (
+    ImageClassificationDataModule,
+)
 from project.utils.env_vars import DATA_DIR, NETWORK_DIR, NUM_WORKERS
 from project.utils.typing_utils import C, H, W
 
@@ -36,7 +38,7 @@ def imagenet_normalization():
 ImageIndex = NewType("ImageIndex", int)
 
 
-class ImageNetDataModule(VisionDataModule):
+class ImageNetDataModule(ImageClassificationDataModule):
     """ImageNet datamodule.
 
     Extracted from https://github.com/Lightning-Universe/lightning-bolts/blob/master/src/pl_bolts/datamodules/imagenet_datamodule.py

project/datamodules/image_classification/imagenet32.py

@@ -17,7 +17,9 @@
 from torchvision.datasets import VisionDataset
 from torchvision.transforms import v2 as transforms
 
-from project.datamodules.vision import VisionDataModule
+from project.datamodules.image_classification.image_classification import (
+    ImageClassificationDataModule,
+)
 from project.utils.env_vars import DATA_DIR, SCRATCH
 from project.utils.typing_utils import C, H, W
 
@@ -167,7 +169,7 @@ def _load_dataset(self):
         self._data_loaded = True
 
 
-class ImageNet32DataModule(VisionDataModule):
+class ImageNet32DataModule(ImageClassificationDataModule):
     """TODO: Add a `val_split` argument, that supports a value of `0`."""
 
     name: ClassVar[str] = "imagenet32"

project/datamodules/image_classification/inaturalist.py

@@ -9,9 +9,7 @@
 import torchvision.transforms as T
 from torchvision.datasets import INaturalist
 
-from project.datamodules.image_classification.image_classification import (
-    ImageClassificationDataModule,
-)
+from project.datamodules.vision import VisionDataModule
 from project.utils.env_vars import DATA_DIR, NUM_WORKERS, SLURM_TMPDIR
 from project.utils.typing_utils import C, H, W
 
@@ -34,7 +32,7 @@ def inat_dataset_dir() -> Path:
     return network_dir
 
 
-class INaturalistDataModule(ImageClassificationDataModule):
+class INaturalistDataModule(VisionDataModule):
     name: ClassVar[str] = "inaturalist"
     """Dataset name."""
 

project/utils/hydra_config_utils.py

@@ -171,15 +171,23 @@ def get_all_configs_in_group_of_type(
             )
         }
 
-    return [
-        name
-        for name, object_type in names_to_types.items()
-        if (
-            issubclass(object_type, config_target_type)
-            if include_subclasses
-            else object_type in config_target_type
-        )
-    ]
+    def _matches_protocol(object: type, protocol: type) -> bool:
+        return isinstance(object, protocol)  # todo: weird!
+
+    compatible_config_names = []
+    for name, object_type in names_to_types.items():
+        if not include_subclasses:
+            if object_type in config_target_type:
+                compatible_config_names.append(name)
+            continue
+        for t in config_target_type:
+            if (
+                issubclass(t, typing.Protocol) and _matches_protocol(object_type, t)
+            ) or issubclass(object_type, t):
+                compatible_config_names.append(name)
+                break
+
+    return compatible_config_names
 
 
 def get_all_configs_in_group_with_target(group_name: str, some_type: type) -> list[str]: