Introduce AdaNet controller.

This AdaNet controller performs a two phase ensembling of deeper and deeper neural network architectures.

PiperOrigin-RevId: 284378509

adanet/experimental/controllers/BUILD

@@ -30,3 +30,19 @@ py_library(
         "//adanet/experimental/work_units:work_unit",
     ],
 )
+
+py_library(
+    name = "adanet_controller",
+    srcs = ["adanet_controller.py"],
+    srcs_version = "PY3",
+    visibility = ["//adanet/experimental:__subpackages__"],
+    deps = [
+        ":controller",
+        "//adanet/experimental/keras:ensemble_model",
+        "//adanet/experimental/phases:phase",
+        "//adanet/experimental/storages:in_memory_storage",
+        "//adanet/experimental/storages:storage",
+        "//adanet/experimental/work_units:keras_trainer",
+        "//adanet/experimental/work_units:work_unit",
+    ],
+)

adanet/experimental/controllers/adanet_controller.py

@@ -0,0 +1,151 @@
+# Lint as: python3
+# Copyright 2019 The AdaNet Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""An AdaNet controller for model search."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import itertools
+
+from typing import Callable, Iterator, List, Sequence, Union
+from adanet.experimental.controllers.controller import Controller
+from adanet.experimental.keras.ensemble_model import MeanEnsemble
+from adanet.experimental.phases.phase import Phase
+from adanet.experimental.storages.in_memory_storage import InMemoryStorage
+from adanet.experimental.storages.storage import Storage
+from adanet.experimental.work_units.keras_trainer import KerasTrainer
+from adanet.experimental.work_units.work_unit import WorkUnit
+
+import tensorflow as tf
+
+
+class AdaNetCandidatePhase(Phase):
+  """Generates and trains neural networks with various layer depths."""
+
+  def __init__(self, dataset: tf.data.Dataset,
+               candidates_per_iteration: int,
+               optimizer: Union[str, tf.keras.optimizers.Optimizer],
+               loss: Union[str, tf.keras.losses.Loss],
+               output_units: int,
+               metrics: List[Union[str, tf.keras.metrics.Metric]] = None,
+               units_per_layer: int = 128,
+               layer_activation: Union[str, Callable[..., tf.Tensor]] = 'relu',
+               output_activation: Union[str, Callable[...,
+                                                      tf.Tensor]] = 'linear'):
+    self._dataset = dataset
+    self._candidates_per_iteration = candidates_per_iteration
+    self._optimizer = optimizer
+    self._loss = loss
+    self._metrics = metrics
+    self._units_per_layer = units_per_layer
+    self._output_units = output_units
+    self._layer_activation = layer_activation
+    self._output_activation = output_activation
+    self._candidate_storage = None
+
+  # TODO: Add warning about build not being called.
+  def build(self, candidate_storage: Storage):
+    self._candidate_storage = candidate_storage
+
+  def work_units(self) -> Iterator[WorkUnit]:
+    for network in self._generate_networks():
+      yield KerasTrainer(network, self._dataset, self._candidate_storage)
+
+  def _generate_networks(self) -> Iterator[tf.keras.Model]:
+    best_candidate = self._candidate_storage.get_best_models(num_models=1)
+    if not best_candidate:
+      num_layers = 0
+    else:
+      num_layers = len(best_candidate[0].layers)
+    for i in range(self._candidates_per_iteration):
+      model = tf.keras.Sequential()
+      for _ in range(num_layers+i):
+        model.add(tf.keras.layers.Dense(units=self._units_per_layer,
+                                        activation=self._layer_activation))
+      model.add(tf.keras.layers.Dense(units=self._output_units,
+                                      activation=self._output_activation))
+      model.compile(optimizer=self._optimizer,
+                    loss=self._loss,
+                    metrics=self._metrics)
+      yield model
+
+
+# TODO: Make this a more general phase.
+class AdaNetEnsemblePhase(Phase):
+  """Ensembles submodels."""
+
+  def __init__(self, dataset: tf.data.Dataset,
+               candidates_per_iteration: int,
+               optimizer: Union[str, tf.keras.optimizers.Optimizer],
+               loss: Union[str, tf.keras.losses.Loss],
+               metrics: List[Union[str, tf.keras.metrics.Metric]] = None):
+    self._dataset = dataset
+    self._candidates_per_iteration = candidates_per_iteration
+    self._optimizer = optimizer
+    self._loss = loss
+    self._metrics = metrics
+    self._candidate_storage = None
+    self._ensemble_storage = None
+
+  def build(self, candidate_storage: Storage, ensemble_storage: Storage):
+    self._candidate_storage = candidate_storage
+    self._ensemble_storage = ensemble_storage
+
+  @property
+  def ensemble_storage(self):
+    return self._ensemble_storage
+
+  # TODO: Revisit how newest candidates are obtained within this
+  # phase.
+  def work_units(self) -> Iterator[WorkUnit]:
+    best_candidates = self._candidate_storage.get_newest_models(
+        num_models=self._candidates_per_iteration)
+    best_ensemble = self._ensemble_storage.get_best_models(num_models=1)
+    for candidate in best_candidates:
+      if not best_ensemble:
+        ensemble = MeanEnsemble([candidate])
+      else:
+        ensemble = MeanEnsemble(best_ensemble[0].submodels + [candidate])
+      ensemble.compile(optimizer=self._optimizer,
+                       loss=self._loss,
+                       metrics=self._metrics)
+
+      yield KerasTrainer(ensemble, self._dataset, self._ensemble_storage)
+
+
+class AdaNetController(Controller):
+  """A controller that trains candidate networks and ensembles them."""
+
+  def __init__(self,
+               candidate_phase: AdaNetCandidatePhase,
+               ensemble_phase: AdaNetEnsemblePhase,
+               iterations: int,
+               candidate_storage: Storage = InMemoryStorage(),
+               ensemble_storage: Storage = InMemoryStorage()):
+    candidate_phase.build(candidate_storage)
+    ensemble_phase.build(candidate_storage, ensemble_storage)
+    self._candidate_phase = candidate_phase
+    self._ensemble_phase = ensemble_phase
+    self._iterations = iterations
+
+  def work_units(self) -> Iterator[WorkUnit]:
+    for _ in range(self._iterations):
+      for work_unit in itertools.chain(self._candidate_phase.work_units(),
+                                       self._ensemble_phase.work_units()):
+        yield work_unit
+
+  def get_best_models(self, num_models) -> Sequence[tf.keras.Model]:
+    return self._ensemble_phase.ensemble_storage.get_best_models(num_models)

adanet/experimental/keras/BUILD

@@ -33,6 +33,7 @@ py_strict_test(
         ":ensemble_model",
         ":model_search",
         ":testing_utils",
+        "//adanet/experimental/controllers:adanet_controller",
         "//adanet/experimental/controllers:sequential_controller",
         "//adanet/experimental/phases:keras_tuner_phase",
         "//adanet/experimental/phases:train_keras_models_phase",

adanet/experimental/keras/ensemble_model.py

@@ -27,14 +27,19 @@
 class EnsembleModel(tf.keras.Model):
   """An ensemble of Keras models."""
 
-  def __init__(self, submodels: Sequence[tf.keras.Model]):
+  def __init__(self, submodels: Sequence[tf.keras.Model],
+               freeze_submodels: bool = True):
     """Initializes an EnsembleModel.
 
     Args:
       submodels: A list of `tf.keras.Model` that compose the ensemble.
+      freeze_submodels: Whether to freeze the weights of submodels.
     """
 
     super().__init__()
+    if freeze_submodels:
+      for submodel in submodels:
+        submodel.trainable = False
     self._submodels = submodels
 
   @property
@@ -49,6 +54,9 @@ class MeanEnsemble(EnsembleModel):
   """An ensemble that averages submodel outputs."""
 
   def call(self, inputs):
+    if len(self._submodels) == 1:
+      return self._submodels[0](inputs)
+
     submodel_outputs = []
     for submodel in self.submodels:
       submodel_outputs.append(submodel(inputs))

adanet/experimental/keras/model_search_test.py

@@ -24,6 +24,9 @@
 
 from absl import flags
 from absl.testing import parameterized
+from adanet.experimental.controllers.adanet_controller import AdaNetCandidatePhase
+from adanet.experimental.controllers.adanet_controller import AdaNetController
+from adanet.experimental.controllers.adanet_controller import AdaNetEnsemblePhase
 from adanet.experimental.controllers.sequential_controller import SequentialController
 from adanet.experimental.keras import testing_utils
 from adanet.experimental.keras.ensemble_model import MeanEnsemble
@@ -77,7 +80,10 @@ def test_phases_end_to_end(self):
     model2.compile(
         optimizer=tf.keras.optimizers.Adam(0.01), loss='mse', metrics=['mae'])
 
-    ensemble = MeanEnsemble(submodels=[model1, model2])
+    # TODO: This test could potentially have the best model be
+    # a non-ensemble Keras model. Therefore, need to address this issue and
+    # remove the freeze_submodels flag.
+    ensemble = MeanEnsemble(submodels=[model1, model2], freeze_submodels=False)
     ensemble.compile(
         optimizer=tf.keras.optimizers.Adam(0.01), loss='mse', metrics=['mae'])
 
@@ -89,13 +95,6 @@ def test_phases_end_to_end(self):
         TrainKerasModelsPhase([ensemble], dataset=train_dataset),
     ])
 
-    train_dataset, test_dataset = testing_utils.get_test_data(
-        train_samples=128,
-        test_samples=64,
-        input_shape=(10,),
-        num_classes=2,
-        random_seed=42)
-
     model_search = ModelSearch(controller)
     model_search.run()
     self.assertIsInstance(
@@ -158,6 +157,41 @@ def build_ensemble():
     self.assertIsInstance(
         model_search.get_best_models(num_models=1)[0], MeanEnsemble)
 
+  def test_adanet_controller_end_to_end(self):
+    train_dataset, test_dataset = testing_utils.get_test_data(
+        train_samples=1280,
+        test_samples=640,
+        input_shape=(10,),
+        num_classes=10,
+        random_seed=42)
+
+    train_dataset = train_dataset.batch(32)
+    test_dataset = test_dataset.batch(32)
+
+    candidate_phase = AdaNetCandidatePhase(
+        train_dataset,
+        candidates_per_iteration=2,
+        optimizer='adam',
+        loss='sparse_categorical_crossentropy',
+        output_units=10)
+    # TODO: Setting candidates_per_iteration greater than the one
+    # for the candidate phase will lead to unexpected behavior.
+    ensemble_phase = AdaNetEnsemblePhase(
+        train_dataset,
+        candidates_per_iteration=2,
+        optimizer='adam',
+        loss='sparse_categorical_crossentropy')
+
+    adanet_controller = AdaNetController(
+        candidate_phase,
+        ensemble_phase,
+        iterations=5)
+
+    model_search = ModelSearch(adanet_controller)
+    model_search.run()
+    self.assertIsInstance(
+        model_search.get_best_models(num_models=1)[0], MeanEnsemble)
+
 
 if __name__ == '__main__':
   tf.enable_v2_behavior()

adanet/experimental/phases/phase.py

@@ -31,6 +31,7 @@ class Phase(abc.ABC):
   A phase is only complete once all its work units complete, as a barrier.
   """
 
+  # TODO: Remove this build function.
   def build(self, storage: Storage, previous: 'Phase' = None):
     self._storage = storage
     self._previous = previous

adanet/experimental/storages/in_memory_storage.py

@@ -52,3 +52,9 @@ def load_model(self, model_id: int) -> tf.keras.Model:
 
   def get_best_models(self, num_models) -> Sequence[tf.keras.Model]:
     return [m for _, _, m in heapq.nsmallest(num_models, self._models)]
+
+  def get_newest_models(self, num_models) -> Sequence[tf.keras.Model]:
+    return [
+        m for _, m_id, m in self._models
+        if m_id in [self._id - i for i in range(num_models)]
+    ]

adanet/experimental/storages/storage.py

@@ -42,3 +42,7 @@ def get_best_models(self, num_models) -> Sequence[tf.keras.Model]:
     # TODO: Rethink get_best_model API since it's defined in Storage,
     # Phases, and Controllers.
     pass
+
+  @abc.abstractmethod
+  def get_newest_models(self, num_models) -> Sequence[tf.keras.Model]:
+    pass

adanet/experimental/work_units/keras_trainer.py

@@ -36,4 +36,10 @@ def __init__(self, model: tf.keras.Model, dataset: tf.data.Dataset,
   def execute(self):
     self._model.fit(self._dataset)
     results = self._model.evaluate(self._dataset)
-    self._storage.save_model(self._model, results[0])
+    # If the model was compiled with metrics, the results is a list of loss +
+    # metric values. If the model was compiled without metrics, it is a loss
+    # scalar.
+    if isinstance(results, list):
+      self._storage.save_model(self._model, results[0])
+    else:
+      self._storage.save_model(self._model, results)