Training at scale

Datasets on disk, made by many processes at once, and the training sets of the model built on them; see Training at scale on a SLURM cluster for the SLURM pipeline that uses them.

Shards

Datasets on disk, in shards made by any number of processes.

The training sets of the surrogates grow to millions of waveforms, more than a process makes in a walltime or holds in memory. A ShardedStore keeps one in a directory:

config.json            what the items are (a frozen dataclass with
                       n_binaries, shard_size and seed)
shards/00000.npz       items 0 ... shard_size - 1, and so on
shards/00000.json      when, where and how fast it was made
locks/00000.lock       a shard being made: by whom, refreshed while it is

Each shard is a function of the configuration and of its index alone (its items drawn from the seed sequence (seed, index)), so shards can be made in any order, by any process on any machine sharing the filesystem: a process claims one by creating its lock file exclusively (taking over a lock not refreshed for a while, whose owner died), and writes it atomically. A shard lost is made again, identically. The first N items are a uniform sample for every N, so the training sets of a learning curve are nested.

Subclasses say what a shard holds (ShardedStore._make()) and may keep derived shards next to the base ones, under other source names, which replace some fields (ShardedStore.overlay_fields) of the base shards.

The processes making them run as batch jobs with a walltime: on SIGTERM or SIGUSR1 (which SLURM sends before the walltime) they abandon the shard they are making (stop_on_signals()) and exit with EXIT_REQUEUE, for the job to be run again.

class ShardedStore(directory: str)

A dataset on disk, in shards; see the module docstring. Open one with Store(directory), make one with create().

Subclasses set config_class (a frozen dataclass with n_binaries, shard_size and seed, to_json and from_json) and implement _make(); those with derived shards also source_directory() and overlay_fields.

chunks(n: int | None = None, fields: Sequence[str] | None = None, source: str = 'base', select: Callable[[ndarray], ndarray] | None = None) → Iterator[Dict[str, ndarray]]

The fields (all, by default; and index, the items’ positions in the dataset) of its first n valid items (all, by default), a shard at a time; of those whose indices pass select (a mask) if given. The overlay_fields are those of source.

classmethod create(directory: str, config)

Make the dataset config in directory, or open it if it is there: the same, or with fewer items and a full last shard (then it is grown).

generate(n_jobs: int = -1, max_seconds: float | None = None, stale_seconds: float = 600.0, **kwargs) → bool

Make the shards not yet there (see work_through()), each with n_jobs processes; whether all are there.

load(n: int | None = None, fields: Sequence[str] | None = None, source: str = 'base', select=None) → Dict[str, ndarray]

chunks(), concatenated.

overlay_fields: Sequence[str] = ()

Fields of the base shards which those of the other sources replace.

path(index: int, source: str = 'base') → str

The shard index of source.

seed(index: int) → list

The seed (sequence) of the items of shard index.

source_directory(source: str) → str

Where the shards of source are ("base": the directory).

status(source: str = 'base') → str

How far the shards of source are, as text.

work_through(source: str, work: Callable[[int], None], max_seconds: float | None = None, stale_seconds: float = 600.0, poll_seconds: float = 30.0) → bool

Claim and work(index) the shards of source not yet there, until there are none (waiting for those others are making), or max_seconds have passed: whether all are there.

class ShardLock(filename: str, stale_seconds: float)

An exclusive claim on a shard, a lock file refreshed by a thread while held; one not refreshed for stale_seconds was left by a process that died, and is taken over. Two processes taking over the same stale lock both make the shard, identically.

write_atomically(filename: str, write: Callable) → None

write(file) to a temporary file next to filename, synced, then renamed over it: a reader sees all of it or none.

save_arrays(filename: str, **arrays) → None

np.savez (uncompressed), atomically.

stop_on_signals() → None

Raise Stopped on SIGTERM and SIGUSR1: what is being made is abandoned (its lock released), what is made is kept.

EXIT_REQUEUE = 75

The exit status of a command stopped before it finished, to be run again.

The modes model

Training sets of the modes surrogate (Model) at scale.

Model.generate makes a training set in one process and holds it in memory; ShardedModesDataset keeps one on disk, in shards, made by any number of processes on any number of machines (see mlgw_bns.sharding), and trains models on the first N waveforms of it a shard at a time. A dataset is a directory:

config.json         what the waveforms are (ModesDatasetConfig)
downsampling.npz    the downsampling indices of every mode
shards/00000.npz    waveforms 0 ... shard_size - 1, and so on
pca/<n>.npz         the principal components of the first n waveforms

A shard holds, for each waveform, its parameters (the five of ParameterSet.parameter_array) and the amplitude (single precision) and phase residuals of every mode at its downsampling indices, as Model._multimode_mode_residuals gives them from one EOB call: about 36 kB a waveform for the seven modes of hom7_big, and a second of TEOBResumS. The downsampling indices are trained first, once (ShardedModesDataset.train_downsampling()), and must be the same for a validation set and the training set it validates (ShardedModesDataset.copy_downsampling()).

Training a model (ShardedModesDataset.train()) fits the principal components on the first pca_size waveforms (once, ShardedModesDataset.principal_components()), reduces the residuals of the first N shard by shard, and fits a regressor of each mode on them (ModeModel.fit_nn), saving each as it is done: a training interrupted resumes from the modes it had finished, and from the checkpoint of the one it was training if that is a JaxMLPNetwork.

class ModesDatasetConfig(n_binaries: int, shard_size: int = 1024, seed: int = 1, modes: tuple = ((2, 2), (2, 1), (3, 1), (3, 2), (3, 3), (4, 3), (4, 4)), initial_frequency_hz: float = 5.0, srate_hz: float = 4096.0, reference_amplitude: bool = True, parameter_ranges: ParameterRanges = <factory>)

What the waveforms of a ShardedModesDataset are, and the Model they train.

Parameters:

• n_binaries (int) – How many; a dataset whose last shard is full can be grown.

• shard_size (int, optional) – Waveforms a shard.

  Default: 1024

• seed (int, optional) – The parameters of shard k are drawn from the seed sequence (seed, k): datasets with different seeds are independent.

  Default: 1

• modes (optional) – Of the Model; the defaults are those of hom7_big (make_hom7_dataset.py).

  Default: ((2, 2), (2, 1), (3, 1), (3, 2), (3, 3), (4, 3), (4, 4))

• initial_frequency_hz (optional) – Of the Model; the defaults are those of hom7_big (make_hom7_dataset.py).

  Default: 5.0

• srate_hz (optional) – Of the Model; the defaults are those of hom7_big (make_hom7_dataset.py).

  Default: 4096.0

• reference_amplitude (optional) – Of the Model; the defaults are those of hom7_big (make_hom7_dataset.py).

  Default: True

• parameter_ranges (optional) – Of the Model; the defaults are those of hom7_big (make_hom7_dataset.py).

  Default: <factory>

model(filename: str | None = None, **kwargs) → Model

A Model of these settings (not trained); kwargs of its modes’ ModeModel, such as nn_kind.

same_waveforms(other: ModesDatasetConfig) → bool

Whether other has waveforms of the same model (if not the same ones): what a validation set and its training set share.

class ShardedModesDataset(directory: str)

A training (or validation) set of a Model on disk; see the module docstring. Open one with ShardedModesDataset(directory), make one with create().

config_class

alias of ModesDatasetConfig

copy_downsampling(other: ShardedModesDataset) → None

Use the downsampling indices of other (of waveforms of the same model): those of the training set, for a validation set.

downsampling() → dict

mode -> DownsamplingIndices.

generate(n_jobs: int = -1, max_seconds: float | None = None, stale_seconds: float = 600.0) → bool

Make the shards not yet there (see work_through()), each with n_jobs processes; whether all are there. The downsampling indices must be there.

load_residuals(n: int | None = None) → tuple

The parameters (n, 5) of the first n valid waveforms and mode -> Residuals of every mode.

model(filename: str | None = None, **kwargs) → Model

A Model of this dataset (see ModesDatasetConfig.model()), with its downsampling indices if they are there.

parameters(index: int) → list

The WaveformParameters of shard index.

principal_components(n: int, n_components: int = 30) → dict

mode -> PrincipalComponentData of the first n valid waveforms, fitted once (as ModeModel.generate does) and kept in pca/<n>.npz.

reduced(model: Model, mode: Mode, n: int) → tuple

The parameters (n, 5) of the first n valid waveforms, the principal components of their residuals of mode in the model’s basis, and their power weights (None for a mode not weighted; see ModeModel._training_power_weights), reduced a shard at a time.

residuals(mode: Mode, n: int | None = None, select: Callable | None = None) → Iterator[tuple]

(parameters, Residuals) of mode for the first n valid waveforms (all, by default), a shard at a time.

train(n: int, filename: str, nn_kind: type | ~typing.Mapping[~mlgw_bns.pn_modes.Mode, type] = <class 'mlgw_bns.neural_network.KernelRidgeNetwork'>, hyperparameters: ~typing.Callable[[~mlgw_bns.pn_modes.Mode, int], ~mlgw_bns.neural_network.Hyperparameters] | None = None, pca_size: int | None = None, modes: ~typing.Sequence | None = None, checkpoint: str | None = None, sizes: ~typing.Mapping[~mlgw_bns.pn_modes.Mode, int] | None = None) → Model | None

Train a Model on the first n valid waveforms and save it as filename (without its training data, like the packaged one).

Parameters:

• n (int) – Training waveforms.

• filename (str) – Base filename of the model (its files are {filename}_l2_m2_nn.pkl and so on).

• nn_kind (type or mapping, optional) – Of the regressors, or of the regressor of each mode (kernel ridge for those not in it).

  Default: <class 'mlgw_bns.neural_network.KernelRidgeNetwork'>

• hyperparameters (callable, optional) – (mode, n) -> Hyperparameters; by default those of ModeModel.default_hyperparameters.

  Default: None

• pca_size (int, optional) – The principal components are those of the first pca_size waveforms (principal_components()); n by default.

  Default: None

• modes (sequence, optional) – Train only the regressors of these modes (each saved as it is done, skipped if it is there): for trainings split in jobs. The model is saved once those of all the modes are there.

  Default: None

• checkpoint (str, optional) – Base filename of the checkpoints of the training of each mode ({checkpoint}_l2_m2.joblib), for the regressors that keep one (JaxMLPNetwork).

  Default: None

• sizes (mapping, optional) – Train the regressors of these modes on the first this many waveforms instead of n.

  Default: None

Returns:

The model, once all of it is trained.

Return type:

Model, optional

train_downsampling(size: int = 64, n_jobs: int = 1) → None

Train the downsampling indices of every mode on size waveforms each (Model.train_downsampling), unless they are there.

load_model(dataset: ShardedModesDataset, filename: str) → Model

The model saved as filename by ShardedModesDataset.train().

stored_size(model: Model) → dict

Bytes on disk of the regressors and of the rest (downsampling, principal components, metadata) of a saved model, and their sum.