-
Notifications
You must be signed in to change notification settings - Fork 142
Enable building an ensemble model from the cross validation checkpoints of a BYO Lightning model #529
Enable building an ensemble model from the cross validation checkpoints of a BYO Lightning model #529
Changes from all commits
1250de0
efc34fb
f94b2e9
7430785
229b7ea
16a09c3
ee957d8
596efee
2c2160f
fe6fa93
45895e7
41f1b48
f6bb7a2
3e3b069
c0de1e6
0889a88
f98b22e
d73f769
f4dfbe7
2767e18
8a63e0a
116e566
7107e27
d0c7724
6735d2c
585511b
4df8c09
05a67dc
8d93907
e7ba7f1
e4ebfcb
7625a48
f5b288c
a79673f
3e9ec7e
559a717
f7be221
a3179a8
f7bdc7f
939ec5a
f206057
d71f498
6987f99
f6e025f
500bf11
0d97fe9
b204f08
dd17d78
388e0a8
d727fe1
fa7a6e4
581c6a9
1134c0f
e35db5b
9093b7e
bf072c0
5a76cc1
2d75d24
e064483
168eb29
838bb48
0f3690c
9c7f6b4
d612c6e
6861178
98c7683
54c845e
48aca37
d3d8477
a4ea25a
5e7f0b1
a71fa90
0874a74
88bb96d
424b2e7
4903b96
547ed7d
538a7a4
002886c
c6ca034
484fe4a
b771137
6d57d08
6a9fd75
cea51e4
9fc696a
b947452
f7d22f4
af079ae
38b12b7
c255774
2db1857
21044e6
4f10169
e0e31e6
1d23735
1486bff
ff8fef1
ddb3c9f
e71d49b
db18fdd
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -19,8 +19,8 @@ | |
from InnerEye.Common.fixed_paths import DEFAULT_AML_UPLOAD_DIR, DEFAULT_LOGS_DIR_NAME | ||
from InnerEye.Common.generic_parsing import GenericConfig | ||
from InnerEye.Common.type_annotations import PathOrString, TupleFloat2 | ||
from InnerEye.ML.common import DATASET_CSV_FILE_NAME, ModelExecutionMode, create_unique_timestamp_id, \ | ||
get_best_checkpoint_path, get_recovery_checkpoint_path | ||
from InnerEye.ML.common import (DATASET_CSV_FILE_NAME, ModelExecutionMode, create_unique_timestamp_id, | ||
get_best_checkpoint_path, get_recovery_checkpoint_path) | ||
|
||
# A folder inside of the outputs folder that will contain all information for running the model in inference mode | ||
|
||
|
@@ -198,6 +198,8 @@ class WorkflowParams(param.Parameterized): | |
cross_validation_split_index: int = param.Integer(DEFAULT_CROSS_VALIDATION_SPLIT_INDEX, bounds=(-1, None), | ||
doc="The index of the cross validation fold this model is " | ||
"associated with when performing k-fold cross validation") | ||
ensemble_model_name: str = param.String( | ||
doc=("The class name of the ensemble model to build from the cross validation checkpoints of a Lightning model.")) | ||
Comment on lines
+201
to
+202
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I feel that this extra requirement is a key bottleneck in your design. Anybody who wants to implement an ensemble model now needs to define the model itself (ensemble member), and then another class that is the ensemble logic. For the other models, we have been able to achieve that without extra overhead - for example, we can define a Prostate model, and get Prostate ensemble models for free. |
||
inference_on_train_set: Optional[bool] = \ | ||
param.Boolean(None, | ||
doc="If set, enable/disable full image inference on training set after training.") | ||
|
@@ -829,7 +831,7 @@ def load_checkpoint_and_modify(self, path_to_checkpoint: Path) -> Dict[str, Any] | |
|
||
def load_checkpoint(path_to_checkpoint: Path, use_gpu: bool = True) -> Dict[str, Any]: | ||
""" | ||
Loads a Torch checkpoint from the given file. If use_gpu==False, map all parameters to the GPU, otherwise | ||
Loads a Torch checkpoint from the given file. If use_gpu==False, map all parameters to the CPU, otherwise | ||
left the device of all parameters unchanged. | ||
""" | ||
import torch | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -3,23 +3,25 @@ | |
# Licensed under the MIT License (MIT). See LICENSE in the repo root for license information. | ||
# ------------------------------------------------------------------------------------------ | ||
import abc | ||
from typing import Any, Dict, Iterator, List, Optional, Tuple | ||
import logging | ||
from copy import deepcopy | ||
from pathlib import Path | ||
from typing import Any, Dict, Iterator, List, Optional, Tuple | ||
|
||
import param | ||
import torch | ||
from azureml.core import ScriptRunConfig | ||
from azureml.train.hyperdrive import GridParameterSampling, HyperDriveConfig, PrimaryMetricGoal, choice | ||
from pytorch_lightning import LightningDataModule, LightningModule | ||
from torch.optim import Optimizer | ||
from torch.optim.lr_scheduler import _LRScheduler | ||
from azureml.core import ScriptRunConfig | ||
from azureml.train.hyperdrive import GridParameterSampling, HyperDriveConfig, PrimaryMetricGoal, choice | ||
|
||
from InnerEye.Azure.azure_util import CROSS_VALIDATION_SPLIT_INDEX_TAG_KEY | ||
from InnerEye.Common.generic_parsing import GenericConfig, create_from_matching_params | ||
from InnerEye.Common.metrics_constants import TrackedMetrics | ||
from InnerEye.ML.common import ModelExecutionMode | ||
from InnerEye.ML.deep_learning_config import DatasetParams, OptimizerParams, OutputParams, TrainerParams, \ | ||
WorkflowParams | ||
from InnerEye.ML.deep_learning_config import (DatasetParams, OptimizerParams, OutputParams, TrainerParams, | ||
WorkflowParams, load_checkpoint) | ||
from InnerEye.ML.utils import model_util | ||
from InnerEye.ML.utils.lr_scheduler import SchedulerWithWarmUp | ||
from InnerEye.ML.utils.run_recovery import RunRecovery | ||
|
@@ -35,11 +37,11 @@ class InnerEyeInference(abc.ABC): | |
|
||
model.on_inference_start() | ||
for dataset_split in [Train, Val, Test] | ||
model.on_inference_epoch_start(dataset_split, is_ensemble_model=False) | ||
for batch_idx, item in enumerate(dataloader[dataset_split])): | ||
model_outputs = model.forward(item) | ||
model.inference_step(item, batch_idx, model_outputs) | ||
model.on_inference_epoch_end() | ||
model.on_inference_start_dataset(dataset_split, is_ensemble_model=False) | ||
for batch_idx, batch in enumerate(dataloader[dataset_split])): | ||
posteriors = model.forward(batch) | ||
model.record_posteriors(batch, batch_idx, posteriors) | ||
model.on_inference_end_dataset() | ||
model.on_inference_end() | ||
""" | ||
|
||
|
@@ -50,29 +52,28 @@ def on_inference_start(self) -> None: | |
""" | ||
pass | ||
|
||
def on_inference_epoch_start(self, dataset_split: ModelExecutionMode, is_ensemble_model: bool) -> None: | ||
def on_inference_start_dataset(self, execution_mode: ModelExecutionMode) -> None: | ||
""" | ||
Runs initialization for inference, when starting inference on a new dataset split (train/val/test). | ||
Depending on the settings, this can be called anywhere between 0 (no inference at all) to 3 times (inference | ||
on all of train/val/test split). | ||
:param dataset_split: Indicates whether the item comes from the training, validation or test set. | ||
:param is_ensemble_model: If False, the model_outputs come from an individual model. If True, the model | ||
outputs come from multiple models. | ||
:param execution_mode: Indicates whether the item comes from the training, validation or test set. | ||
""" | ||
pass | ||
|
||
def inference_step(self, batch: Any, batch_idx: int, model_output: torch.Tensor) -> None: | ||
def record_posteriors(self, batch: Any, batch_idx: int, posteriors: torch.Tensor) -> None: | ||
""" | ||
This hook is called when the model has finished making a prediction. It can write the results to a file, | ||
or compute metrics and store them. | ||
:param batch: The batch of data for which the model made a prediction. | ||
:param model_output: The model outputs. This would usually be a torch.Tensor, but can be any datatype. | ||
:param batch_idx: The index of the batch. | ||
:param posteriors: The posteriors output by the model. | ||
""" | ||
# We don't want abstract methods here, it avoids class creation for unit tests, and we also want this | ||
# method to be left optional (it should be possible to also use Lightning's native test_step method) | ||
raise NotImplementedError("Method on_inference_start must be overwritten in a derived class.") | ||
|
||
def on_inference_epoch_end(self) -> None: | ||
def on_inference_end_dataset(self) -> None: | ||
""" | ||
Called when the inference on one of the dataset splits (train/val/test) has finished. | ||
Depending on the settings, this can be called anywhere between 0 (no inference at all) to 3 times (inference | ||
|
@@ -87,10 +88,137 @@ def on_inference_end(self) -> None: | |
""" | ||
pass | ||
|
||
def aggregate_ensemble_model_outputs(self, model_outputs: Iterator[torch.Tensor]) -> torch.Tensor: | ||
|
||
class InnerEyeEnsembleInference(): | ||
""" | ||
InnerEyeEnsembleInference provides help for buiding ensemble models from cross validation runs of LightningModules, | ||
and doing inference on the ensemble. | ||
|
||
To set up an ensemble and then do inference, call | ||
model.load_checkpoints_into_ensemble(checkpoints, use_gpu) | ||
model.on_ensemble_inference_start() | ||
for dataset_split in [Train, Val, Test] | ||
model.on_ensemble_inference_start_dataset(dataset_split, is_ensemble_model=False) | ||
for batch_idx, batch in enumerate(dataloader[dataset_split])): | ||
posterior = model.ensemble_forward(batch) | ||
model.record_ensemble_posterior(batch, batch_idx, posterior) | ||
model.on_ensemble_inference_end_dataset() | ||
model.on_ensemble_inference_end() | ||
|
||
We have not duplicated the method documentation from InnerEyeInference, where you can find further explantaion of | ||
the role of each method. | ||
""" | ||
def __init__(self, outputs_folder: Optional[Path] = None) -> None: | ||
""" | ||
Sets up the list of models that forms the ensemble. These models should inherit from both InnerEyeEinference and | ||
LightiningModule, but since mypy does not have support for specifying interesection types we specify just one, | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't get this? |
||
InnerEyeInference, here and assert the other when needed. | ||
:param outputs_folder: The root directory where we store metric files and reports. This is marked optional so | ||
that subclasses may be instantiated without constructor arguments. | ||
""" | ||
self.outputs_folder = outputs_folder | ||
self.ensemble_models: List[LightningModule] = [] | ||
|
||
def load_checkpoints_into_ensemble( # type: ignore | ||
self, | ||
exemplar: LightningModule, | ||
checkpoint_paths: List[Path], | ||
use_gpu: bool) -> None: | ||
""" | ||
Convenience method takes each checkpoint path from a list of checkpoint paths and adds it as an additional | ||
member of the ensemble. | ||
:param exemplar: Each model in the ensemble will be based on a copy of this exemplar. | ||
:param checkpoint_paths: A list of paths to checkpoints for loading into new models in the ensemble. | ||
:param use_gpu: Passed on eventaully to deep_learning_config.load_checkpoint. | ||
""" | ||
for checkpoint_path in checkpoint_paths: | ||
self._load_checkpoint_into_ensemble(exemplar, checkpoint_path, use_gpu) | ||
|
||
def _load_checkpoint_into_ensemble( # type: ignore | ||
self, | ||
exemplar: LightningModule, | ||
checkpoint_path: Path, | ||
use_gpu: bool) -> None: | ||
""" | ||
Load a single checkpoint path as an additional member of the ensemble. | ||
:param exemplar: Each model in the ensemble will be based on a copy of this exemplar. | ||
:param checkpoint_path: The path to the to checkpoint file to load into a new model in the ensemble. | ||
:param use_gpu: Passed on to deep_learning_config.load_checkpoint. | ||
""" | ||
checkpoint = load_checkpoint(checkpoint_path, use_gpu) | ||
# new_model = deepcopy(exemplar) # In AML, this line triggers "ReferenceError: weakly-referenced object no | ||
# longer exists" so we will make a new instance instead | ||
logging.info(f"Adding a {type(exemplar)} to the ensemble and loading its state from a checkpoint.") | ||
new_model = type(exemplar)() | ||
assert isinstance(new_model, LightningModule) # mypy | ||
new_model.load_state_dict(checkpoint['state_dict'], strict=False) | ||
self.ensemble_models.append(new_model) | ||
|
||
def on_ensemble_inference_start(self) -> None: | ||
""" | ||
Set each model in the enseble into evaluation mode. | ||
""" | ||
for model in self.ensemble_models: | ||
assert isinstance(model, LightningModule) # mypy | ||
model.eval() | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. calling model.eval() is critical. If users forget to call the superclass method, they will run their model in a wrong way. Can we bake that into local_checkpoint? |
||
|
||
def on_ensemble_inference_start_dataset(self, execution_mode: ModelExecutionMode) -> None: | ||
""" | ||
Runs initialization for inference, when starting inference on a new dataset split (train/val/test). | ||
:param execution_mode: Passed on to InnerEyeInference model in the ensemble. | ||
""" | ||
pass | ||
|
||
def ensemble_forward(self, batch_x: torch.Tensor) -> torch.Tensor: | ||
""" | ||
Aggregate the posteriors from each model in the ensemble using our static method | ||
`aggregate_ensemble_model_outputs` | ||
:param batch_x: The batch of data for which the model made a prediction. | ||
:returns: The aggregated ensemble outputs. | ||
""" | ||
model_outputs: List[torch.Tensor] = [] | ||
for model in self.ensemble_models: | ||
assert isinstance(model, LightningModule) # mypy | ||
model_outputs.append(model.forward(batch_x)) | ||
posterior = InnerEyeEnsembleInference.aggregate_ensemble_model_outputs(iter(model_outputs)) | ||
return posterior | ||
|
||
def record_ensemble_posterior(self, batch_y: torch.Tensor, batch_idx: int, posterior: torch.Tensor) -> None: | ||
""" | ||
This hook is called when the model has finished making a prediction, and can write the results to a file, or | ||
compute metrics and store them. In this base class we cannot know which metrics are needed so you will need to | ||
override this method. Here are two ideas: | ||
|
||
self.test_mse.append(torch.nn.functional.mse_loss(posterior, batch["y"])) | ||
self.test_mae.update(preds=posterior, target=batch["y"]) | ||
|
||
(where `self.test_mse: List[torch.Tensor]` and `self.test_mae: MeanAbsoluteError` would be initialised in the | ||
constructor). | ||
|
||
:param batch_y: The batch of data for which the model made a prediction. | ||
:param batch_idx: The index of the batch | ||
:param posterior: The posterior, typically from ensemble_forward. | ||
""" | ||
pass | ||
|
||
def on_ensemble_inference_end_dataset(self) -> None: | ||
""" | ||
Run after the inference run over a test/val/train dataset | ||
""" | ||
pass | ||
|
||
def on_ensemble_inference_end(self) -> None: | ||
""" | ||
Run after all the test/val/train datasets' inference runs. Override to save metrics to disk and generate reports. | ||
""" | ||
pass | ||
|
||
@staticmethod | ||
def aggregate_ensemble_model_outputs(model_outputs: Iterator[torch.Tensor]) -> torch.Tensor: | ||
""" | ||
Aggregates the outputs of multiple models when using an ensemble model. In the default implementation, | ||
this averages the tensors coming from all the models. | ||
Aggregates the outputs of multiple models when using an ensemble model. In the default implementation, this | ||
example averages the tensors coming from all of the models in the ensemble. Override this static method if you | ||
need a different aggregation from the mean we use here. | ||
:param model_outputs: An iterator over the model outputs for all ensemble members. | ||
:return: The aggregate model outputs. | ||
""" | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not sure I get the question? This method here should only return a freshly generated model, the _model properties are populated elsewhere.