OnnxExporter

OnnxExporter(game_configuration: GameConfiguration, model_configuration: ModelConfiguration, training_configuration: TrainingConfiguration, create_module_and_observation_space: Callable[[int | None], tuple[torch.nn.Module, Box]], export_lock: Callable[[], contextlib.AbstractContextManager]) -> None

Exports a checkpointed model to ONNX.

export is verbose, for the standalone Experiment.export() path. export_quiet additionally holds the caller-provided export lock and suppresses noisy third-party logging, for use from the background checkpoint-write path, which may run concurrently with learner updates.

Initialize the exporter.

Parameters:

• game_configuration (GameConfiguration) – The experiment's game configuration, embedded as ONNX metadata.
• model_configuration (ModelConfiguration) – The experiment's model configuration, embedded as ONNX metadata.
• training_configuration (TrainingConfiguration) – The experiment's training configuration, embedded as ONNX metadata.
• create_module_and_observation_space (Callable[[int | None], tuple[Module, Box]]) – Builds the exportable module and its observation space for a given checkpoint number (or the live model when None).
• export_lock (Callable[[], AbstractContextManager]) – The lock to hold during a quiet (background) export.

Source

Methods

export

export(output_path: str | PathLike, *, checkpoint_number: int | None = None) -> Path

Export the model to ONNX, logging progress.

Parameters:

• output_path (str | PathLike) – Where to write the ONNX file.
• checkpoint_number (int | None) – The checkpoint to export, or None to export the live model.

Returns:

• Path – The output path.

Source

export_quiet

export_quiet(output_path: str | PathLike, *, checkpoint_number: int | None) -> None

Export the model to ONNX while holding the export lock, suppressing routine logging.

Parameters:

• output_path (str | PathLike) – Where to write the ONNX file.
• checkpoint_number (int | None) – The checkpoint to export.

Source