Merge branch 'selective-model-load' of https://github.com/Ayaro/auto-sklearn into Ayaro-selective-model-load

Author: mfeurer

autosklearn/automl.py

@@ -639,11 +639,16 @@ def _load_models(self):
         else:
             seed = self._seed
 
-        self.models_ = self._backend.load_all_models(seed)
+        self.ensemble_ = self._backend.load_ensemble(seed)
+        if self.ensemble_:
+            identifiers = self.ensemble_.identifiers_
+            self.models_ = self._backend.load_models_by_identifiers(identifiers)
+        else:
+            self.models_ = self._backend.load_all_models(seed)
+
         if len(self.models_) == 0:
             raise ValueError('No models fitted!')
 
-        self.ensemble_ = self._backend.load_ensemble(seed)
 
     def score(self, X, y):
         # fix: Consider only index 1 of second dimension

autosklearn/util/backend.py

@@ -168,19 +168,44 @@ def load_all_models(self, seed):
             model_files = os.listdir(model_directory)
             model_files = [os.path.join(model_directory, mf) for mf in model_files]
 
+        models = self.load_models_by_file_names(model_files)
+
+        return models
+
+    def load_models_by_file_names(self, model_file_names):
         models = dict()
-        for model_file in model_files:
+
+        for model_file in model_file_names:
             # File names are like: {seed}.{index}.model
             if model_file.endswith('/'):
                 model_file = model_file[:-1]
             basename = os.path.basename(model_file)
-            automl_seed = int(basename.split('.')[0])
-            idx = int(basename.split('.')[1])
-            with open(os.path.join(model_directory, basename), 'rb') as fh:
-                models[(automl_seed, idx)] = (pickle.load(fh))
+
+            basename_parts = basename.split('.')
+            seed = int(basename_parts[0])
+            idx = int(basename_parts[1])
+
+            models[(seed, idx)] = self.load_model_by_seed_and_id(seed, idx)
+
+        return models
+
+    def load_models_by_identifiers(self, identifiers):
+        models = dict()
+
+        for identifier in identifiers:
+            seed, idx = identifier
+            models[identifier] = self.load_model_by_seed_and_id(seed, idx)
 
         return models
 
+    def load_model_by_seed_and_id(self, seed, idx):
+        model_directory = self.get_model_dir()
+        model_file_name = '%s.%s.model' % (seed, idx)
+        model_file_path = os.path.join(model_directory, model_file_name)
+
+        with open(model_file_path, 'rb') as fh:
+            return (pickle.load(fh))
+
     def get_ensemble_dir(self):
         return os.path.join(self.internals_directory, 'ensembles')
 

test/automl/test_models.py

@@ -0,0 +1,49 @@
+# -*- encoding: utf-8 -*-
+from __future__ import print_function
+import unittest
+import mock
+from autosklearn.automl import AutoML
+from autosklearn.util.backend import Backend
+
+
+class AutoMLStub(AutoML):
+
+    def __init__(self):
+        self.__class__ = AutoML
+
+
+class AutoMlModelsTest(unittest.TestCase):
+
+    def setUp(self):
+        self.automl = AutoMLStub()
+        self.automl._shared_mode = False
+        self.automl._seed = 42
+        self.automl._backend = mock.Mock(spec=Backend)
+        self.automl._delete_output_directories = lambda: 0
+
+    def test_only_loads_ensemble_models(self):
+        identifiers = [(1, 2), (3, 4)]
+        models = [ 42 ]
+        self.automl._backend.load_ensemble.return_value.identifiers_ \
+            = identifiers
+        self.automl._backend.load_models_by_identifiers.side_effect \
+            = lambda ids: models if ids is identifiers else None
+
+        self.automl._load_models()
+
+        self.assertEqual(models, self.automl.models_)
+
+    def test_loads_all_models_if_no_ensemble(self):
+        models = [ 42 ]
+        self.automl._backend.load_ensemble.return_value = None
+        self.automl._backend.load_all_models.return_value = models
+
+        self.automl._load_models()
+
+        self.assertEqual(models, self.automl.models_)
+
+    def test_raises_if_no_models(self):
+        self.automl._backend.load_ensemble.return_value = None
+        self.automl._backend.load_all_models.return_value = []
+
+        self.assertRaises(ValueError, self.automl._load_models)

test/util/test_backend.py

@@ -0,0 +1,61 @@
+# -*- encoding: utf-8 -*-
+from __future__ import print_function
+import unittest
+import mock
+from autosklearn.util.backend import Backend
+
+from sys import version_info
+if version_info.major == 2:
+    import __builtin__ as builtins
+else:
+    import builtins
+
+
+class BackendModelsTest(unittest.TestCase):
+
+    class BackendStub(Backend):
+
+        def __init__(self):
+            self.__class__ = Backend
+
+    def setUp(self):
+        self.model_directory = '/model_directory/'
+        self.backend = self.BackendStub()
+        self.backend.get_model_dir = lambda: self.model_directory
+
+    @mock.patch('six.moves.cPickle.load')
+    @mock.patch.object(builtins, 'open')
+    def test_loads_model_by_seed_and_id(self, openMock, pickleLoadMock):
+        seed = 13
+        idx = 17
+        expected_model = self._setup_load_model_mocks(openMock, pickleLoadMock, seed, idx)
+
+        actual_model = self.backend.load_model_by_seed_and_id(seed, idx)
+
+        self.assertEqual(expected_model, actual_model)
+
+    @mock.patch('six.moves.cPickle.load')
+    @mock.patch.object(builtins, 'open')
+    def test_loads_models_by_identifiers(self, openMock, pickleLoadMock):
+        seed = 13
+        idx = 17
+        expected_model = self._setup_load_model_mocks(openMock, pickleLoadMock, seed, idx)
+        expected_dict = { (seed, idx): expected_model }
+
+        actual_dict = self.backend.load_models_by_identifiers([(seed, idx)])
+
+        self.assertIsInstance(actual_dict, dict)
+        self.assertDictEqual(expected_dict, actual_dict)
+
+    def _setup_load_model_mocks(self, openMock, pickleLoadMock, seed, idx):
+        model_path = '/model_directory/%s.%s.model' % (seed, idx)
+        file_handler = 'file_handler'
+        expected_model = 'model'
+
+        fileMock = mock.MagicMock()
+        fileMock.__enter__.return_value = file_handler
+
+        openMock.side_effect = lambda path, flag: fileMock if path == model_path and flag == 'rb' else None
+        pickleLoadMock.side_effect = lambda fh: expected_model if fh == file_handler else None
+
+        return expected_model