def generate_and_print_model_summary(config: ModelConfigBase, model: DeviceAwareModule) -> None:
"""
Writes a human readable summary of the present model to logging.info, and logs the number of trainable
parameters to AzureML.
:param config: The configuration for the model.
:param model: The instantiated Pytorch model.
"""
random_state = RandomStateSnapshot.snapshot_random_state()
# There appears to be a bug in apex, where previous use (in training for example) causes problems
# when another model is later built on the CPU (for example, before loading from a checkpoint)
# https://github.com/NVIDIA/apex/issues/694
# Hence, move the model to the GPU before doing model summary.
if config.use_gpu:
model = model.cuda()
if isinstance(config, ScalarModelBase):
# To generate the model summary, read the first item of the dataset. Then use the model's own
# get_model_input function to convert the dataset item to input tensors, and feed them through the model.
train_dataset = config.get_torch_dataset_for_inference(ModelExecutionMode.TRAIN)
train_item_0 = next(iter(train_dataset.as_data_loader(shuffle=False, batch_size=1, num_dataload_workers=0)))
model_inputs = get_scalar_model_inputs_and_labels(config, model, train_item_0).model_inputs
# The model inputs may already be converted to float16, assuming that we would do mixed precision.
# However, the model is not yet converted to float16 when this function is called, hence convert back to float32
summary = ModelSummary(model)
summary.generate_summary(input_tensors=model_inputs, log_summaries_to_files=config.log_summaries_to_files)
elif config.is_segmentation_model:
summary_for_segmentation_models(config, model)
assert model.summarizer
summary = model.summarizer # type: ignore
else:
raise ValueError("Don't know how to generate a summary for this type of model?")
RUN_CONTEXT.log(LoggingColumns.NumTrainableParameters, summary.n_trainable_params)
random_state.restore_random_state()
def log_to_azure(self, label: str, metric: float) -> None:
"""
Logs a metric as a key/value pair to AzureML.
"""
if not is_offline_run_context(RUN_CONTEXT):
metric_name = self.logging_prefix + label
RUN_CONTEXT.log(metric_name, metric)
# When running in a cross validation setting, log all metrics to the hyperdrive parent run too,
# so that we can easily overlay graphs across runs.
if self.log_to_parent_run and PARENT_RUN_CONTEXT:
if self.cross_validation_split_index > DEFAULT_CROSS_VALIDATION_SPLIT_INDEX:
PARENT_RUN_CONTEXT.log(f"{metric_name}_Split{self.cross_validation_split_index}",
metric)
def run(self) -> None:
"""
Driver function to run a ML experiment. If an offline cross validation run is requested, then
this function is recursively called for each cross validation split.
"""
if self.is_offline_cross_val_parent_run():
if self.model_config.is_segmentation_model:
raise NotImplementedError(
"Offline cross validation is only supported for classification models."
)
self.spawn_offline_cross_val_classification_child_runs()
return
# Get the AzureML context in which the script is running
if not self.model_config.is_offline_run and PARENT_RUN_CONTEXT is not None:
logging.info("Setting tags from parent run.")
self.set_run_tags_from_parent()
self.save_build_info_for_dotnet_consumers()
# Set data loader start method
self.set_multiprocessing_start_method()
# configure recovery container if provided
checkpoint_handler = CheckpointHandler(model_config=self.model_config,
azure_config=self.azure_config,
project_root=self.project_root,
run_context=RUN_CONTEXT)
checkpoint_handler.discover_and_download_checkpoints_from_previous_runs(
)
# do training and inference, unless the "only register" switch is set (which requires a run_recovery
# to be valid).
if not self.azure_config.register_model_only_for_epoch:
# Set local_dataset to the mounted path specified in azure_runner.py, if any, or download it if that fails
# and config.local_dataset was not already set.
self.model_config.local_dataset = self.mount_or_download_dataset()
self.model_config.write_args_file()
logging.info(str(self.model_config))
# Ensure that training runs are fully reproducible - setting random seeds alone is not enough!
make_pytorch_reproducible()
# Check for existing dataset.csv file in the correct locations. Skip that if a dataset has already been
# loaded (typically only during tests)
if self.model_config.dataset_data_frame is None:
assert self.model_config.local_dataset is not None
ml_util.validate_dataset_paths(self.model_config.local_dataset)
# train a new model if required
if self.azure_config.train:
with logging_section("Model training"):
model_train(self.model_config, checkpoint_handler)
else:
self.model_config.write_dataset_files()
self.create_activation_maps()
# log the number of epochs used for model training
RUN_CONTEXT.log(name="Train epochs",
value=self.model_config.num_epochs)
# We specify the ModelProcessing as DEFAULT here even if the run_recovery points to an ensemble run, because
# the current run is a single one. See the documentation of ModelProcessing for more details.
best_epoch = self.run_inference_and_register_model(
checkpoint_handler, ModelProcessing.DEFAULT)
# Generate report
if best_epoch:
Runner.generate_report(self.model_config, best_epoch,
ModelProcessing.DEFAULT)
elif self.model_config.is_scalar_model and len(
self.model_config.get_test_epochs()) == 1:
# We don't register scalar models but still want to create a report if we have run inference.
Runner.generate_report(self.model_config,
self.model_config.get_test_epochs()[0],
ModelProcessing.DEFAULT)
def model_train(config: ModelConfigBase,
checkpoint_handler: CheckpointHandler) -> ModelTrainingResults:
"""
The main training loop. It creates the model, dataset, optimizer_type, and criterion, then proceeds
to train the model. If a checkpoint was specified, then it loads the checkpoint before resuming training.
:param config: The arguments which specify all required information.
:param checkpoint_handler: Checkpoint handler object to find checkpoint paths for model initialization
:raises TypeError: If the arguments are of the wrong type.
:raises ValueError: When there are issues loading a previous checkpoint.
"""
# Save the dataset files for later use in cross validation analysis
config.write_dataset_files()
# set the random seed for all libraries
ml_util.set_random_seed(config.get_effective_random_seed(),
"Patch visualization")
# Visualize how patches are sampled for segmentation models. This changes the random generator, but we don't
# want training to depend on how many patients we visualized, and hence set the random seed again right after.
with logging_section(
"Visualizing the effect of sampling random crops for training"):
visualize_random_crops_for_dataset(config)
ml_util.set_random_seed(config.get_effective_random_seed(),
"Model training")
logging.debug("Creating the PyTorch model.")
# Create the train loader and validation loader to load images from the dataset
data_loaders = config.create_data_loaders()
# Get the path to the checkpoint to recover from
checkpoint_path = checkpoint_handler.get_recovery_path_train()
models_and_optimizer = ModelAndInfo(
config=config,
model_execution_mode=ModelExecutionMode.TRAIN,
checkpoint_path=checkpoint_path)
# Create the main model
# If continuing from a previous run at a specific epoch, then load the previous model.
model_loaded = models_and_optimizer.try_create_model_and_load_from_checkpoint(
)
if not model_loaded:
raise ValueError(
"There was no checkpoint file available for the model for given start_epoch {}"
.format(config.start_epoch))
# Print out a detailed breakdown of layers, memory consumption and time.
generate_and_print_model_summary(config, models_and_optimizer.model)
# Move model to GPU and adjust for multiple GPUs
models_and_optimizer.adjust_model_for_gpus()
# Create the mean teacher model and move to GPU
if config.compute_mean_teacher_model:
mean_teacher_model_loaded = models_and_optimizer.try_create_mean_teacher_model_load_from_checkpoint_and_adjust(
)
if not mean_teacher_model_loaded:
raise ValueError(
"There was no checkpoint file available for the mean teacher model "
f"for given start_epoch {config.start_epoch}")
# Create optimizer
models_and_optimizer.create_optimizer()
if checkpoint_handler.should_load_optimizer_checkpoint():
optimizer_loaded = models_and_optimizer.try_load_checkpoint_for_optimizer(
)
if not optimizer_loaded:
raise ValueError(
f"There was no checkpoint file available for the optimizer for given start_epoch "
f"{config.start_epoch}")
# Create checkpoint directory for this run if it doesn't already exist
logging.info(f"Models are saved at {config.checkpoint_folder}")
if not config.checkpoint_folder.is_dir():
config.checkpoint_folder.mkdir()
# Create the SummaryWriters for Tensorboard
writers = create_summary_writers(config)
config.create_dataframe_loggers()
# Create LR scheduler
l_rate_scheduler = SchedulerWithWarmUp(config,
models_and_optimizer.optimizer)
# Training loop
logging.info("Starting training")
train_results_per_epoch, val_results_per_epoch, learning_rates_per_epoch = [], [], []
resource_monitor = None
if config.monitoring_interval_seconds > 0:
# initialize and start GPU monitoring
diagnostics_events = config.logs_folder / "diagnostics"
logging.info(
f"Starting resource monitor, outputting to {diagnostics_events}")
resource_monitor = ResourceMonitor(
interval_seconds=config.monitoring_interval_seconds,
tensorboard_folder=diagnostics_events)
resource_monitor.start()
gradient_scaler = GradScaler(
) if config.use_gpu and config.use_mixed_precision else None
optimal_temperature_scale_values = []
for epoch in config.get_train_epochs():
logging.info("Starting epoch {}".format(epoch))
save_epoch = config.should_save_epoch(
epoch) and models_and_optimizer.optimizer is not None
# store the learning rates used for each epoch
epoch_lrs = l_rate_scheduler.get_last_lr()
learning_rates_per_epoch.append(epoch_lrs)
train_val_params: TrainValidateParameters = \
TrainValidateParameters(data_loader=data_loaders[ModelExecutionMode.TRAIN],
model=models_and_optimizer.model,
mean_teacher_model=models_and_optimizer.mean_teacher_model,
epoch=epoch,
optimizer=models_and_optimizer.optimizer,
gradient_scaler=gradient_scaler,
epoch_learning_rate=epoch_lrs,
summary_writers=writers,
dataframe_loggers=config.metrics_data_frame_loggers,
in_training_mode=True)
training_steps = create_model_training_steps(config, train_val_params)
train_epoch_results = train_or_validate_epoch(training_steps)
train_results_per_epoch.append(train_epoch_results.metrics)
metrics.validate_and_store_model_parameters(writers.train, epoch,
models_and_optimizer.model)
# Run without adjusting weights on the validation set
train_val_params.in_training_mode = False
train_val_params.data_loader = data_loaders[ModelExecutionMode.VAL]
# if temperature scaling is enabled then do not save validation metrics for the checkpoint epochs
# as these will be re-computed after performing temperature scaling on the validation set.
if isinstance(config, SequenceModelBase):
train_val_params.save_metrics = not (
save_epoch and config.temperature_scaling_config)
training_steps = create_model_training_steps(config, train_val_params)
val_epoch_results = train_or_validate_epoch(training_steps)
val_results_per_epoch.append(val_epoch_results.metrics)
if config.is_segmentation_model:
metrics.store_epoch_stats_for_segmentation(
config.outputs_folder, epoch, epoch_lrs,
train_epoch_results.metrics, val_epoch_results.metrics)
if save_epoch:
# perform temperature scaling if required
if isinstance(
config,
SequenceModelBase) and config.temperature_scaling_config:
optimal_temperature, scaled_val_results = \
temperature_scaling_steps(config, train_val_params, val_epoch_results)
optimal_temperature_scale_values.append(optimal_temperature)
# overwrite the metrics for the epoch with the metrics from the temperature scaled model
val_results_per_epoch[-1] = scaled_val_results.metrics
models_and_optimizer.save_checkpoint(epoch)
# Updating the learning rate should happen at the end of the training loop, so that the
# initial learning rate will be used for the very first epoch.
l_rate_scheduler.step()
model_training_results = ModelTrainingResults(
train_results_per_epoch=train_results_per_epoch,
val_results_per_epoch=val_results_per_epoch,
learning_rates_per_epoch=learning_rates_per_epoch,
optimal_temperature_scale_values_per_checkpoint_epoch=
optimal_temperature_scale_values)
logging.info("Finished training")
# Since we have trained the model further, let the checkpoint_handler object know so it can handle
# checkpoints correctly.
checkpoint_handler.additional_training_done()
# Upload visualization directory to AML run context to be able to see it
# in the Azure UI.
if config.max_batch_grad_cam > 0 and config.visualization_folder.exists():
RUN_CONTEXT.upload_folder(name=VISUALIZATION_FOLDER,
path=str(config.visualization_folder))
writers.close_all()
config.metrics_data_frame_loggers.close_all()
if resource_monitor:
# stop the resource monitoring process
logging.info(
"Shutting down the resource monitor process. Aggregate resource utilization:"
)
for name, value in resource_monitor.read_aggregate_metrics():
logging.info(f"{name}: {value}")
if not is_offline_run_context(RUN_CONTEXT):
RUN_CONTEXT.log(name, value)
resource_monitor.kill()
return model_training_results
def run(self) -> None:
"""
Driver function to run a ML experiment. If an offline cross validation run is requested, then
this function is recursively called for each cross validation split.
"""
if self.is_offline_cross_val_parent_run():
if self.model_config.is_segmentation_model:
raise NotImplementedError("Offline cross validation is only supported for classification models.")
self.spawn_offline_cross_val_classification_child_runs()
return
# Get the AzureML context in which the script is running
if not self.model_config.is_offline_run and PARENT_RUN_CONTEXT is not None:
logging.info("Setting tags from parent run.")
self.set_run_tags_from_parent()
self.save_build_info_for_dotnet_consumers()
# Set data loader start method
self.set_multiprocessing_start_method()
# configure recovery container if provided
checkpoint_handler = CheckpointHandler(model_config=self.model_config,
azure_config=self.azure_config,
project_root=self.project_root,
run_context=RUN_CONTEXT)
checkpoint_handler.download_recovery_checkpoints_or_weights()
# do training and inference, unless the "only register" switch is set (which requires a run_recovery
# to be valid).
if not self.azure_config.only_register_model:
# Set local_dataset to the mounted path specified in azure_runner.py, if any, or download it if that fails
# and config.local_dataset was not already set.
self.model_config.local_dataset = self.mount_or_download_dataset()
# Check for existing dataset.csv file in the correct locations. Skip that if a dataset has already been
# loaded (typically only during tests)
if self.model_config.dataset_data_frame is None:
assert self.model_config.local_dataset is not None
ml_util.validate_dataset_paths(
self.model_config.local_dataset,
self.model_config.dataset_csv)
# train a new model if required
if self.azure_config.train:
with logging_section("Model training"):
model_train(self.model_config, checkpoint_handler, num_nodes=self.azure_config.num_nodes)
else:
self.model_config.write_dataset_files()
self.create_activation_maps()
# log the number of epochs used for model training
RUN_CONTEXT.log(name="Train epochs", value=self.model_config.num_epochs)
# We specify the ModelProcessing as DEFAULT here even if the run_recovery points to an ensemble run, because
# the current run is a single one. See the documentation of ModelProcessing for more details.
self.run_inference_and_register_model(checkpoint_handler, ModelProcessing.DEFAULT)
if self.model_config.generate_report:
self.generate_report(ModelProcessing.DEFAULT)
# If this is an cross validation run, and the present run is child run 0, then wait for the sibling runs,
# build the ensemble model, and write a report for that.
if self.model_config.number_of_cross_validation_splits > 0:
if self.model_config.should_wait_for_other_cross_val_child_runs():
self.wait_for_runs_to_finish()
self.create_ensemble_model()
def log_metrics(self,
metrics: Dict[str, float],
step: Optional[int] = None) -> None:
if self.is_azureml_run:
for key, value in metrics.items():
RUN_CONTEXT.log(key, value)
def model_train(config: ModelConfigBase,
checkpoint_handler: CheckpointHandler,
num_nodes: int = 1) -> ModelTrainingResults:
"""
The main training loop. It creates the Pytorch model based on the configuration options passed in,
creates a Pytorch Lightning trainer, and trains the model.
If a checkpoint was specified, then it loads the checkpoint before resuming training.
:param config: The arguments which specify all required information.
:param checkpoint_handler: Checkpoint handler object to find checkpoint paths for model initialization
:param num_nodes: The number of nodes to use in distributed training.
"""
# Get the path to the checkpoint to recover from
checkpoint_path = checkpoint_handler.get_recovery_path_train()
# This reads the dataset file, and possibly sets required pre-processing objects, like one-hot encoder
# for categorical features, that need to be available before creating the model.
config.read_dataset_if_needed()
# Create the trainer object. Backup the environment variables before doing that, in case we need to run a second
# training in the unit tests.d
old_environ = dict(os.environ)
seed_everything(config.get_effective_random_seed())
trainer, storing_logger = create_lightning_trainer(config, checkpoint_path, num_nodes=num_nodes)
logging.info(f"GLOBAL_RANK: {os.getenv('GLOBAL_RANK')}, LOCAL_RANK {os.getenv('LOCAL_RANK')}. "
f"trainer.global_rank: {trainer.global_rank}")
logging.debug("Creating the PyTorch model.")
lightning_model = create_lightning_model(config)
lightning_model.storing_logger = storing_logger
resource_monitor = None
# Execute some bookkeeping tasks only once if running distributed:
if is_rank_zero():
config.write_args_file()
logging.info(str(config))
# Save the dataset files for later use in cross validation analysis
config.write_dataset_files()
logging.info(f"Model checkpoints are saved at {config.checkpoint_folder}")
# set the random seed for all libraries
ml_util.set_random_seed(config.get_effective_random_seed(), "Patch visualization")
# Visualize how patches are sampled for segmentation models. This changes the random generator, but we don't
# want training to depend on how many patients we visualized, and hence set the random seed again right after.
with logging_section("Visualizing the effect of sampling random crops for training"):
visualize_random_crops_for_dataset(config)
# Print out a detailed breakdown of layers, memory consumption and time.
generate_and_print_model_summary(config, lightning_model.model)
if config.monitoring_interval_seconds > 0:
# initialize and start GPU monitoring
diagnostics_events = config.logs_folder / "diagnostics"
logging.info(f"Starting resource monitor, outputting to {diagnostics_events}")
resource_monitor = ResourceMonitor(interval_seconds=config.monitoring_interval_seconds,
tensorboard_folder=diagnostics_events)
resource_monitor.start()
# Training loop
logging.info("Starting training")
lightning_data = TrainingAndValidationDataLightning(config) # type: ignore
# When trying to store the config object in the constructor, it does not appear to get stored at all, later
# reference of the object simply fail. Hence, have to set explicitly here.
lightning_data.config = config
trainer.fit(lightning_model, datamodule=lightning_data)
trainer.logger.close() # type: ignore
lightning_model.close_all_loggers()
world_size = getattr(trainer, "world_size", 0)
is_azureml_run = not config.is_offline_run
# Per-subject model outputs for regression models are written per rank, and need to be aggregated here.
# Each thread per rank will come here, and upload its files to the run outputs. Rank 0 will later download them.
if is_azureml_run and world_size > 1 and isinstance(lightning_model, ScalarLightning):
upload_output_file_as_temp(lightning_model.train_subject_outputs_logger.csv_path, config.outputs_folder)
upload_output_file_as_temp(lightning_model.val_subject_outputs_logger.csv_path, config.outputs_folder)
# DDP will start multiple instances of the runner, one for each GPU. Those should terminate here after training.
# We can now use the global_rank of the Lightining model, rather than environment variables, because DDP has set
# all necessary properties.
if lightning_model.global_rank != 0:
logging.info(f"Terminating training thread with rank {lightning_model.global_rank}.")
sys.exit()
logging.info("Choosing the best checkpoint and removing redundant files.")
cleanup_checkpoint_folder(config.checkpoint_folder)
# Lightning modifies a ton of environment variables. If we first run training and then the test suite,
# those environment variables will mislead the training runs in the test suite, and make them crash.
# Hence, restore the original environment after training.
os.environ.clear()
os.environ.update(old_environ)
if world_size and isinstance(lightning_model, ScalarLightning):
if is_azureml_run and world_size > 1:
# In a DDP run on the local box, all ranks will write to local disk, hence no download needed.
# In a multi-node DDP, each rank would upload to AzureML, and rank 0 will now download all results and
# concatenate
for rank in range(world_size):
for mode in [ModelExecutionMode.TRAIN, ModelExecutionMode.VAL]:
file = mode.value + "/" + get_subject_output_file_per_rank(rank)
RUN_CONTEXT.download_file(name=TEMP_PREFIX + file, output_file_path=config.outputs_folder / file)
# Concatenate all temporary file per execution mode
for mode in [ModelExecutionMode.TRAIN, ModelExecutionMode.VAL]:
temp_files = (config.outputs_folder / mode.value).rglob(SUBJECT_OUTPUT_PER_RANK_PREFIX + "*")
result_file = config.outputs_folder / mode.value / SUBJECT_METRICS_FILE_NAME
for i, file in enumerate(temp_files):
temp_file_contents = file.read_text()
if i == 0:
# Copy the first file as-is, including the first line with the column headers
result_file.write_text(temp_file_contents)
else:
# For all files but the first one, cut off the header line.
result_file.write_text(os.linesep.join(temp_file_contents.splitlines()[1:]))
model_training_results = ModelTrainingResults(
train_results_per_epoch=list(storing_logger.to_metrics_dicts(prefix_filter=TRAIN_PREFIX).values()),
val_results_per_epoch=list(storing_logger.to_metrics_dicts(prefix_filter=VALIDATION_PREFIX).values()),
train_diagnostics=lightning_model.train_diagnostics,
val_diagnostics=lightning_model.val_diagnostics,
optimal_temperature_scale_values_per_checkpoint_epoch=[]
)
logging.info("Finished training")
# Since we have trained the model further, let the checkpoint_handler object know so it can handle
# checkpoints correctly.
checkpoint_handler.additional_training_done()
# Upload visualization directory to AML run context to be able to see it
# in the Azure UI.
if config.max_batch_grad_cam > 0 and config.visualization_folder.exists():
RUN_CONTEXT.upload_folder(name=VISUALIZATION_FOLDER, path=str(config.visualization_folder))
if resource_monitor:
# stop the resource monitoring process
logging.info("Shutting down the resource monitor process. Aggregate resource utilization:")
for name, value in resource_monitor.read_aggregate_metrics():
logging.info(f"{name}: {value}")
if not config.is_offline_run:
RUN_CONTEXT.log(name, value)
resource_monitor.kill()
return model_training_results
