gw.training.train_pipeline.initialize_stage

gw.training.train_pipeline.initialize_stage(
    pm,
    wfd,
    stage,
    num_workers,
    resume=False,
)

Initializes training based on PosteriorModel metadata and current stage: * Builds transforms (based on noise settings for current stage); * Builds DataLoaders; * At the beginning of a stage (i.e., if not resuming mid-stage), initializes a new optimizer and scheduler; * Freezes / unfreezes SVD layer of embedding network

Parameters

Name	Type	Description	Default
pm	BasePosteriorModel		required
wfd	WaveformDataset		required
stage	dict	Settings specific to current stage of training	required
num_workers	int		required
resume	bool	Whether training is resuming mid-stage. This controls whether the optimizer and scheduler should be re-initialized based on contents of stage dict.	False

Returns

Name	Type	Description
	(train_loader, test_loader)