fix issue with missing classes

Author: ravinkohli

autoPyTorch/api/base_task.py

@@ -12,7 +12,7 @@
 import unittest.mock
 import warnings
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Union
 
 from ConfigSpace.configuration_space import Configuration, ConfigurationSpace
 
@@ -299,6 +299,7 @@ def _get_dataset_input_validator(
         resampling_strategy: Optional[ResamplingStrategies] = None,
         resampling_strategy_args: Optional[Dict[str, Any]] = None,
         dataset_name: Optional[str] = None,
+        dataset_compression: Optional[Mapping[str, Any]] = None,
     ) -> Tuple[BaseDataset, BaseInputValidator]:
         """
         Returns an object of a child class of `BaseDataset` and
@@ -341,6 +342,7 @@ def get_dataset(
         resampling_strategy: Optional[ResamplingStrategies] = None,
         resampling_strategy_args: Optional[Dict[str, Any]] = None,
         dataset_name: Optional[str] = None,
+        dataset_compression: Optional[Mapping[str, Any]] = None,
     ) -> BaseDataset:
         """
         Returns an object of a child class of `BaseDataset` according to the current task.
@@ -375,7 +377,8 @@ def get_dataset(
             y_test=y_test,
             resampling_strategy=resampling_strategy,
             resampling_strategy_args=resampling_strategy_args,
-            dataset_name=dataset_name)
+            dataset_name=dataset_name,
+            dataset_compression=dataset_compression)
 
         return dataset
 

autoPyTorch/data/utils.py

@@ -240,8 +240,8 @@ def validate_dataset_compression_arg(
                 f"\nmemory_allocation = {memory_allocation}"
                 f"\ndataset_compression = {dataset_compression}"
             )
-        # convert to int so we can directly use
-        dataset_compression["memory_allocation"] = floor(memory_allocation * memory_limit)
+        # convert to required memory so we can directly use
+        dataset_compression["memory_allocation"] = memory_allocation * memory_limit
 
     # "methods" must be non-empty sequence
     if (
@@ -464,7 +464,7 @@ def megabytes(arr: DatasetCompressionInputType) -> float:
 
 def reduce_dataset_size_if_too_large(
     X: DatasetCompressionInputType,
-    memory_allocation: int,
+    memory_allocation: float,
     is_classification: bool,
     random_state: Union[int, np.random.RandomState],
     y: Optional[SupportedTargetTypes] = None,

autoPyTorch/pipeline/components/training/trainer/__init__.py

@@ -18,7 +18,7 @@
 from torch.optim.lr_scheduler import _LRScheduler
 from torch.utils.tensorboard.writer import SummaryWriter
 
-from autoPyTorch.constants import STRING_TO_TASK_TYPES
+from autoPyTorch.constants import CLASSIFICATION_TASKS, STRING_TO_TASK_TYPES
 from autoPyTorch.datasets.base_dataset import BaseDatasetPropertiesType
 from autoPyTorch.pipeline.components.base_choice import autoPyTorchChoice
 from autoPyTorch.pipeline.components.base_component import (
@@ -257,6 +257,14 @@ def _fit(self, X: Dict[str, Any], y: Any = None, **kwargs: Any) -> 'TrainerChoic
         if 'optimize_metric' in X and X['optimize_metric'] not in [m.name for m in metrics]:
             metrics.extend(get_metrics(dataset_properties=X['dataset_properties'], names=[X['optimize_metric']]))
         additional_losses = X['additional_losses'] if 'additional_losses' in X else None
+
+        # Ensure that the split is not missing any class.
+        labels = X['y_train'][X['backend'].load_datamanager().splits[X['split_id']][0]]
+        if STRING_TO_TASK_TYPES[X['dataset_properties']['task_type']] in CLASSIFICATION_TASKS:
+            unique_labels = len(np.unique(labels))
+            if unique_labels < X['dataset_properties']['output_shape']:
+                raise ValueError(f"Expected number of unique labels {unique_labels} in train split: {X['split_id']}"
+                                 f" to be = num_classes {X['dataset_properties']['output_shape']}.")
         self.choice.prepare(
             model=X['network'],
             metrics=metrics,
@@ -268,7 +276,7 @@ def _fit(self, X: Dict[str, Any], y: Any = None, **kwargs: Any) -> 'TrainerChoic
             metrics_during_training=X['metrics_during_training'],
             scheduler=X['lr_scheduler'],
             task_type=STRING_TO_TASK_TYPES[X['dataset_properties']['task_type']],
-            labels=X['y_train'][X['backend'].load_datamanager().splits[X['split_id']][0]],
+            labels=labels,
             step_interval=X['step_interval']
         )
         total_parameter_count, trainable_parameter_count = self.count_parameters(X['network'])