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 _adjust_for_gpus(
cls, model: DeviceAwareModule, config: ModelConfigBase,
model_execution_mode: ModelExecutionMode) -> DeviceAwareModule:
"""
Updates a torch model so that input mini-batches are parallelized across the batch dimension to utilise
multiple gpus. If model parallel is set to True and execution is in test mode, then model is partitioned to
perform full volume inference.
This assumes the model has been created, that the optimizer has not yet been created, and the the model has not
been adjusted twice. This method should not be called externally. Use instead adjust_model_for_gpus
or adjust_mean_teacher_model_for_gpus
:returns Adjusted model
"""
if config.use_gpu:
model = model.cuda()
logging.info(
"Adjusting the model to use mixed precision training.")
# If model parallel is set to True, then partition the network across all available gpus.
if config.use_model_parallel:
devices = config.get_cuda_devices()
assert devices is not None # for mypy
model.partition_model(devices=devices) # type: ignore
else:
logging.info(
"Making no adjustments to the model because no GPU was found.")
# Update model related config attributes (After Model Parallel Activated)
config.adjust_after_mixed_precision_and_parallel(model)
# DataParallel enables running the model with multiple gpus by splitting samples across GPUs
# If the model is used in training mode, data parallel is activated by default.
# Similarly, if model parallel is not activated, data parallel is used as a backup option
use_data_parallel = (model_execution_mode == ModelExecutionMode.TRAIN
) or (not config.use_model_parallel)
if config.use_gpu and use_data_parallel:
logging.info("Adjusting the model to use DataParallel")
# Move all layers to the default GPU before activating data parallel.
# This needs to happen even though we put the model to the GPU at the beginning of the method,
# but we may have spread it across multiple GPUs later.
model = model.cuda()
model = DataParallelModel(model,
device_ids=config.get_cuda_devices())
return model