Device Selection

A module to benchmark CPU and GPU training performance and select the best device for training a deep neural network. The module provides a function to run isolated subprocesses for benchmarking both CPU and GPU, and returns the selected device along with the benchmark results.

class seawrd.device_selection.BenchmarkCache(seawrd_config: dict[str, Any], num_inputs: int, tensorflow_version: str, keras_version: str, gpu_names: tuple[str, ...], cpu_name: str, device_choice: DeviceChoice)

Bases: object

A dataclass to hold the cached benchmark results for CPU and GPU.

It contains the benchmark results for both CPU and GPU, allowing for quick access to previously computed results without the need to rerun the benchmarks. This class is used to encapsulate the cached benchmark results, providing a structured way to access the results for both devices.

cpu_name: str

device_choice: DeviceChoice

gpu_names: tuple[str, ...]

keras_version: str

num_inputs: int

seawrd_config: dict[str, Any]

tensorflow_version: str

to_overhead() → dict[str, Any]

Convert the BenchmarkCache instance into a dictionary suitable for generating a cache key. This does not contain the benchmark results themselves, but only the configuration and device information necessary to uniquely identify the cache entry.

Returns:: A dictionary representation of the BenchmarkCache instance, containing only the configuration and device information necessary for generating a unique cache key.
Return type:: dict[str, Any]

class seawrd.device_selection.BenchmarkWorkerConfig(data_path: str, seawrd_config: dict[str, Any], benchmark_epochs: int, benchmark_repeats: int, warmup_epochs: int)

Bases: object

A lightweight payload for the benchmark worker subprocess.

The payload keeps the serialised SEAWRD configuration as a plain dictionary so the benchmark path can stay decoupled from Keras imports until the worker process reconstructs the validated config.

benchmark_epochs: int

benchmark_repeats: int

data_path: str

classmethod from_config(config: SEAWRDConfig | Mapping[str, Any] | Any, data_path: str | Path, benchmark_epochs: int, benchmark_repeats: int, warmup_epochs: int) → BenchmarkWorkerConfig

Build a benchmark worker payload from a SEAWRD configuration or raw mapping.

Parameters:

config (SEAWRDConfig | Mapping[str, Any] | Any) – The configuration object to convert into a worker payload.
data_path (str | Path) – The path to the training and validation data for the benchmark worker.
benchmark_epochs (int) – The number of epochs to train the model during benchmarking.
benchmark_repeats (int) – The number of times to repeat the benchmark for each device.
warmup_epochs (int) – The number of warmup epochs to run before benchmarking.

Returns:

An instance of BenchmarkWorkerConfig populated with the provided parameters and configuration.

Return type:

BenchmarkWorkerConfig

classmethod from_dict(data: dict[str, Any]) → BenchmarkWorkerConfig

Create a BenchmarkWorkerConfig instance from a dictionary.

Parameters:: data (dict[str, Any]) – A dictionary containing the benchmark worker configuration.
Returns:: An instance of BenchmarkWorkerConfig populated with the data from the dictionary.
Return type:: BenchmarkWorkerConfig

seawrd_config: dict[str, Any]

to_dict() → dict[str, Any]

Convert the worker payload into a JSON-serializable dictionary.

Returns:: A dictionary representation of the BenchmarkWorkerConfig instance, suitable for JSON serialisation.
Return type:: dict[str, Any]

warmup_epochs: int

class seawrd.device_selection.DeviceBenchmarkResult(device: str, times: list[float], median_seconds: float, mean_seconds: float, std_seconds: float, steps_per_second: float, seconds_per_epoch: float, gpu_detected: bool | None = None, gpu_devices: list[str] | None = None)

Bases: object

A dataclass to hold the benchmark results for a specific device (CPU or GPU).

It contains the device name, median, mean, and standard deviation of the training times, as well as the final validation loss and steps per second achieved during the benchmark. This class is used to encapsulate the results the benchmarking process for easy comparison and selection of the best device for training.

device: str

classmethod from_dict(data: dict[str, Any]) → DeviceBenchmarkResult

Create a DeviceBenchmarkResult instance from a dictionary.

Parameters:: data (dict[str, Any]) – A dictionary containing the benchmark results for a device.
Returns:: An instance of DeviceBenchmarkResult populated with the data from the dictionary.
Return type:: DeviceBenchmarkResult

gpu_detected: bool | None = None

gpu_devices: list[str] | None = None

mean_seconds: float

median_seconds: float

seconds_per_epoch: float

std_seconds: float

steps_per_second: float

times: list[float]

to_dict() → dict[str, Any]

Convert the DeviceBenchmarkResult instance into a dictionary.

Returns:: A dictionary representation of the DeviceBenchmarkResult instance, suitable for JSON serialisation.
Return type:: dict[str, Any]

class seawrd.device_selection.DeviceChoice(device: str, reason: str, cpu_result: DeviceBenchmarkResult | None = None, gpu_result: DeviceBenchmarkResult | None = None)

Bases: object

A dataclass to hold the result of the device selection process.

It contains the selected device (CPU or GPU), the reason for the selection, and the benchmark results for both CPU and GPU. This class is used to encapsulate the outcome of the device selection process, providing a clear structured way to access the selected device and the associated benchmark results.

cpu_result: DeviceBenchmarkResult | None = None

device: str

classmethod from_dict(data: dict[str, Any]) → DeviceChoice

Create a DeviceChoice instance from a dictionary.

Parameters:: data (dict[str, Any]) – A dictionary containing the device selection results.
Returns:: An instance of DeviceChoice populated with the data from the dictionary.
Return type:: DeviceChoice

gpu_result: DeviceBenchmarkResult | None = None

reason: str

to_dict() → dict[str, Any]

Convert the DeviceChoice instance into a dictionary.

Returns:: A dictionary representation of the DeviceChoice instance, suitable for JSON serialisation.
Return type:: dict[str, Any]

seawrd.device_selection.choose_training_device(config: Mapping[str, Any], x_train: ndarray, y_train: ndarray, x_val: ndarray, y_val: ndarray, benchmark_epochs: int = 10, benchmark_repeats: int = 3, warmup_epochs: int = 10, force_benchmark: bool = False) → DeviceChoice

Benchmark CPU and GPU in isolated subprocesses and choose a backend for training based on the results.

Parameters:

config (Mapping[str, Any]) – A raw mapping containing the configuration for the model and training.
x_train (np.ndarray) – The training input data.
y_train (np.ndarray) – The training target data.
x_val (np.ndarray) – The validation input data.
y_val (np.ndarray) – The validation target data.
benchmark_epochs (int, optional) – The number of epochs to train the model during benchmarking. Default is 10.
benchmark_repeats (int, optional) – The number of times to repeat the benchmark for each device. Default is 3.
warmup_epochs (int, optional) – The number of warmup epochs to run before benchmarking. Default is 10.
force_benchmark (bool, optional) – Whether to force benchmarking even if a cached result exists. Default is False.

Returns:

A dictionary containing the selected device, the reason for the selection, and the benchmark results for both CPU and GPU.

Return type:

dict[str, Any]