Repository: pykale/pykale
Issue: [Bug] Digits notebook - digits access

# 🐛 Bug

Possible break due to PR #210 changes in digits access API.

## To reproduce

Run https://colab.research.google.com/github/pykale/pykale/blob/main/examples/digits_dann_lightn/tutorial.ipynb#scrollTo=237L9WPwrl2t and there will be an error when training the model

```python
%time trainer.fit(model)
results.update(
    is_validation=True, method_name=cfg.DAN.METHOD, seed=seed, metric_values=trainer.callback_metrics,
)
```

** Stack trace/error message **
```

  | Name              | Type                | Params
----------------------------------------------------------
0 | feat              | SmallCNNFeature     | 312 K 
1 | classifier        | ClassNetSmallImage  | 24.4 K
2 | domain_classifier | DomainNetSmallImage | 128 K 
----------------------------------------------------------
466 K     Trainable params
0         Non-trainable params
466 K     Total params
1.864     Total estimated model params size (MB)
Validation sanity check: 0%
0/2 [00:00<?, ?it/s]
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-30-c69caaf3d927> in <module>()
----> 1 get_ipython().magic('time trainer.fit(model)')
      2 results.update(
      3     is_validation=True, method_name=cfg.DAN.METHOD, seed=seed, metric_values=trainer.callback_metrics,
      4 )

26 frames
<decorator-gen-53> in time(self, line, cell, local_ns)

<timed eval> in <module>()

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight, bias)
    438                             _pair(0), self.dilation, self.groups)
    439         return F.conv2d(input, weight, bias, self.stride,
--> 440                         self.padding, self.dilation, self.groups)
    441 
    442     def forward(self, input: Tensor) -> Tensor:

RuntimeError: Given groups=1, weight of size [64, 3, 5, 5], expected input[150, 1, 32, 32] to have 3 channels, but got 1 channels instead
```



Changed files:
tests/loaddata/test_digits_access.py, kale/loaddata/digits_access.py

diff --git a/kale/loaddata/digits_access.py b/kale/loaddata/digits_access.py
index 758409b..f6db173 100644
--- a/kale/loaddata/digits_access.py
+++ b/kale/loaddata/digits_access.py
@@ -23,16 +23,7 @@ class DigitDataset(Enum):
     SVHN = "SVHN"
 
     @staticmethod
-    def get_access(dataset: "DigitDataset", data_path):
-        """Gets data loaders for digit datasets
-
-        Args:
-            dataset (DigitDataset): dataset name
-            data_path (string): root directory of dataset
-
-        Examples::
-            >>> data_access, num_channel = DigitDataset.get_access(dataset, data_path)
-        """
+    def get_channel_numbers(dataset: "DigitDataset"):
         channel_numbers = {
             DigitDataset.MNIST: 1,
             DigitDataset.MNIST_RGB: 3,
@@ -41,16 +32,36 @@ class DigitDataset(Enum):
             DigitDataset.USPS_RGB: 3,
             DigitDataset.SVHN: 3,
         }
+        return channel_numbers[dataset]
 
+    @staticmethod
+    def get_digit_transform(dataset: "DigitDataset", n_channels):
         transform_names = {
             (DigitDataset.MNIST, 1): "mnist32",
+            (DigitDataset.MNIST, 3): "mnist32rgb",
             (DigitDataset.MNIST_RGB, 3): "mnist32rgb",
             (DigitDataset.MNISTM, 3): "mnistm",
             (DigitDataset.USPS, 1): "usps32",
+            (DigitDataset.USPS, 3): "usps32rgb",
             (DigitDataset.USPS_RGB, 3): "usps32rgb",
             (DigitDataset.SVHN, 3): "svhn",
         }
 
+        return transform_names[(dataset, n_channels)]
+
+    @staticmethod
+    def get_access(dataset: "DigitDataset", data_path, num_channels=None):
+        """Gets data loaders for digit datasets
+
+        Args:
+            dataset (DigitDataset): dataset name
+            data_path (string): root directory of dataset
+            num_channels (int): number of channels, defaults to None
+
+        Examples::
+            >>> data_access, num_channel = DigitDataset.get_access(dataset, data_path)
+        """
+
         factories = {
             DigitDataset.MNIST: MNISTDatasetAccess,
             DigitDataset.MNIST_RGB: MNISTDatasetAccess,
@@ -59,11 +70,9 @@ class DigitDataset(Enum):
             DigitDataset.USPS_RGB: USPSDatasetAccess,
             DigitDataset.SVHN: SVHNDatasetAccess,
         }
-
-        num_channels = channel_numbers[dataset]
-        tf = transform_names[(dataset, num_channels)]
-
-        factories[dataset](data_path, tf)
+        if num_channels is None:
+            num_channels = DigitDataset.get_channel_numbers(dataset)
+        tf = DigitDataset.get_digit_transform(dataset, num_channels)
 
         return factories[dataset](data_path, tf), num_channels
 
@@ -80,9 +89,11 @@ class DigitDataset(Enum):
         Examples::
             >>> source_access, target_access, num_channel = DigitDataset.get_source_target(source, target, data_path)
         """
-        src_access, src_n_channels = DigitDataset.get_access(source, data_path)
-        tgt_access, tgt_n_channels = DigitDataset.get_access(target, data_path)
+        src_n_channels = DigitDataset.get_channel_numbers(source)
+        tgt_n_channels = DigitDataset.get_channel_numbers(target)
         num_channels = max(src_n_channels, tgt_n_channels)
+        src_access, src_n_channels = DigitDataset.get_access(source, data_path, num_channels)
+        tgt_access, tgt_n_channels = DigitDataset.get_access(target, data_path, num_channels)
 
         return src_access, tgt_access, num_channels
 
diff --git a/tests/loaddata/test_digits_access.py b/tests/loaddata/test_digits_access.py
index e215ff9..30c53ec 100644
--- a/tests/loaddata/test_digits_access.py
+++ b/tests/loaddata/test_digits_access.py
@@ -15,7 +15,7 @@ from kale.loaddata.multi_domain import (
 
 SOURCES = ["MNIST", "USPS"]
 TARGETS = ["MNISTM", "SVHN"]
-ALL = ["SVHN", "USPS"]  # SOURCES + TARGETS
+ALL = SOURCES + TARGETS  # ["SVHN", "USPS", "MNISTM"]  # SOURCES + TARGETS
 
 WEIGHT_TYPE = ["natural", "balanced", "preset0"]
 DATASIZE_TYPE = ["max", "source"]