add preprocessed_dtype to determine double or float

Author: ravinkohli

autoPyTorch/api/base_task.py

@@ -270,6 +270,7 @@ def build_pipeline(
         include_components: Optional[Dict[str, Any]] = None,
         exclude_components: Optional[Dict[str, Any]] = None,
         search_space_updates: Optional[HyperparameterSearchSpaceUpdates] = None
+    ) -> BasePipeline:
         """
         Build pipeline according to current task
                 Characteristics of the dataset to guide the pipeline

autoPyTorch/pipeline/components/setup/early_preprocessor/EarlyPreprocessing.py

@@ -42,7 +42,8 @@ def transform(self, X: Dict[str, Any]) -> Dict[str, Any]:
         # We need to also save the preprocess transforms for inference
         X.update({
                  'preprocess_transforms': transforms,
-                 'shape_after_preprocessing': X['X_train'].shape[1:]
+                 'shape_after_preprocessing': X['X_train'].shape[1:],
+                 'preprocessed_dtype': X['X_train'].dtype.name
                  })
         return X
 

autoPyTorch/pipeline/components/setup/early_preprocessor/TimeSeriesEarlyPreProcessing.py

@@ -65,7 +65,8 @@ def transform(self, X: Dict[str, Any]) -> Dict[str, Any]:
         # We need to also save the preprocess transforms for inference
         X.update({
             'preprocess_transforms': transforms,
-            'shape_after_preprocessing': X['X_train'].shape[1:]
+            'shape_after_preprocessing': X['X_train'].shape[1:],
+            'preprocessed_dtype': X['X_train'].dtype.name
             })
         return X
 

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

@@ -447,15 +447,18 @@ def _fit(self, X: Dict[str, Any], y: Any = None, **kwargs: Any) -> 'TrainerChoic
             raise RuntimeError("Budget exhausted without finishing an epoch.")
 
         if self.choice.use_stochastic_weight_averaging and self.choice.swa_updated:
+            use_double = 'float64' in X['preprocessed_dtype']
 
             # update batch norm statistics
-            swa_utils.update_bn(loader=X['train_data_loader'], model=self.choice.swa_model.double())
-
+            swa_model = self.choice.swa_model.double() if use_double else self.choice.swa_model
+            swa_utils.update_bn(loader=X['train_data_loader'], model=swa_model)
             # change model
             update_model_state_dict_from_swa(X['network'], self.choice.swa_model.state_dict())
-            if self.choice.use_snapshot_ensemble and len(self.choice.model_snapshots) > 0:
+            if self.choice.use_snapshot_ensemble:
                 # we update only the last network which pertains to the stochastic weight averaging model
-                swa_utils.update_bn(X['train_data_loader'], self.choice.model_snapshots[-1].double())
+                snapshot_model = self.choice.model_snapshots[-1].double() if use_double else self.choice.model_snapshots[-1]
+                swa_utils.update_bn(X['train_data_loader'], snapshot_model)
+                update_model_state_dict_from_swa(X['network_snapshots'][-1], self.choice.swa_model.state_dict())
 
         # wrap up -- add score if not evaluating every epoch
         if not self.eval_valid_each_epoch(X):

test/test_pipeline/components/training/test_image_data_loader.py

@@ -16,7 +16,6 @@ def test_imageloader_build_transform():
 
     fit_dictionary = dict()
     fit_dictionary['dataset_properties'] = dict()
-    fit_dictionary['dataset_properties']['is_small_preprocess'] = unittest.mock.Mock(())
     fit_dictionary['image_augmenter'] = unittest.mock.Mock()
     fit_dictionary['preprocess_transforms'] = unittest.mock.Mock()
 

test/test_pipeline/components/training/test_training.py

@@ -101,7 +101,7 @@ def test_fit_transform(self):
             'y_train': np.array([0, 1, 0]),
             'train_indices': [0, 1],
             'val_indices': [2],
-            'dataset_properties': {'is_small_preprocess': True},
+            'dataset_properties': {},
             'working_dir': '/tmp',
             'split_id': 0,
             'backend': backend,

test/test_pipeline/test_tabular_classification.py

@@ -205,15 +205,12 @@ def test_pipeline_transform(self, fit_dictionary_tabular, exclude):
         # We expect the transformations to be in the pipeline at anytime for inference
         assert 'preprocess_transforms' in transformed_fit_dictionary_tabular.keys()
 
-    @pytest.mark.parametrize("is_small_preprocess", [True, False])
-    def test_default_configuration(self, fit_dictionary_tabular, is_small_preprocess, exclude):
+    def test_default_configuration(self, fit_dictionary_tabular, exclude):
         """Makes sure that when no config is set, we can trust the
         default configuration from the space"""
 
         fit_dictionary_tabular['epochs'] = 5
 
-        fit_dictionary_tabular['is_small_preprocess'] = is_small_preprocess
-
         pipeline = TabularClassificationPipeline(
             dataset_properties=fit_dictionary_tabular['dataset_properties'],
             exclude=exclude)