def _get_experiment_components(
experiment: IExperiment,
stage: str = None,
device: Device = None,
) -> Tuple[Model, Criterion, Optimizer, Scheduler, Device]:
"""
Inner method for `Experiment` components preparation.
Check available torch device, takes model from the experiment
and creates stage-specified criterion, optimizer, scheduler for it.
Args:
stage: experiment stage name of interest
like "pretrain" / "train" / "finetune" / etc
Returns:
tuple: model, criterion, optimizer,
scheduler and device for a given stage and model
"""
model = experiment.get_model(stage)
criterion = experiment.get_criterion(stage)
optimizer = experiment.get_optimizer(stage, model)
scheduler = experiment.get_scheduler(stage, optimizer)
model, criterion, optimizer, scheduler, device = process_components(
model=model,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
distributed_params=experiment.distributed_params,
device=device,
)
return model, criterion, optimizer, scheduler, device
def _get_experiment_callbacks(
experiment: IExperiment, stage: str,
) -> Dict[str, Callback]:
"""Inner method for `Callbacks` preparation.
Takes callbacks from the Experiment
and filters them for distributed master/worker cases.
Args:
stage (str): stage name of interest,
like "pretrain" / "train" / "finetune" / etc
Returns:
OrderedDict[str, Callback]: Ordered dictionary
with callbacks for current experiment stage.
"""
callbacks = experiment.get_callbacks(stage)
callbacks = utils.filter_callbacks_by_node(callbacks)
callbacks = utils.sort_callbacks_by_order(callbacks)
return callbacks