from __future__ import annotations
from typing import TYPE_CHECKING, Any, Optional, Union
import io
import json
import zipfile
if TYPE_CHECKING:
import torch
[docs]
class NeuralNetworkEnsemblePredictor:
"""Client-side predictor for a downloaded neural network ensemble.
Use :meth:`from_file` to load an artifact created with
:meth:`jcmoptimizer.ActiveLearning.download_neural_network`. Predictions are
evaluated locally with PyTorch and return one prediction per ensemble
member.
"""
def __init__(
self,
metadata: dict[str, Any],
member_model_bytes: bytes,
member_state_bytes: bytes,
device: Optional[Union[str, torch.device]] = None,
) -> None:
"""Create a predictor from loaded artifact metadata and model bytes.
Args:
metadata: Artifact metadata read from ``metadata.json``.
member_model_bytes: Serialized exported PyTorch member model.
member_state_bytes: Serialized stacked ensemble member state.
device: PyTorch device for prediction. If omitted, CUDA is used when
available and CPU otherwise.
"""
if metadata.get("format") != "jcmoptimizer_neural_network_ensemble":
raise ValueError(
f"Unsupported neural network artifact format {metadata.get('format')!r}."
)
if metadata.get("artifact_type") != "torch_export_member_vmap":
raise ValueError(
"Unsupported neural network artifact type "
f"{metadata.get('artifact_type')!r}."
)
try:
import torch
from torch.func import functional_call, vmap
except ImportError as error:
raise ImportError(
"NeuralNetworkEnsemblePredictor requires PyTorch. "
"Install the optional torch dependency to use it."
) from error
self.metadata = metadata
self._torch = torch
self.device = torch.device(
device if device is not None else "cuda" if torch.cuda.is_available() else "cpu"
)
exported_program = torch.export.load(io.BytesIO(member_model_bytes))
self.base_model = exported_program.module().to(self.device)
stacked_state = torch.load(
io.BytesIO(member_state_bytes), map_location=self.device, weights_only=True
)
parameter_names = {name for name, _ in self.base_model.named_parameters()}
buffer_names = {name for name, _ in self.base_model.named_buffers()}
self.params = {
name: value.to(self.device)
for name, value in stacked_state.items()
if name in parameter_names
}
self.buffers = {
name: value.to(self.device)
for name, value in stacked_state.items()
if name in buffer_names
}
self.dtype = self._infer_dtype(stacked_state)
self.num_input = int(metadata["num_input"])
def predict_member(
params: dict[str, torch.Tensor],
buffers: dict[str, torch.Tensor],
x: torch.Tensor,
) -> torch.Tensor:
return functional_call(self.base_model, (params, buffers), (x,))
self._vmap = vmap(predict_member, in_dims=(0, 0, None))
[docs]
@classmethod
def from_file(
cls,
filename: str,
device: Optional[Union[str, torch.device]] = None,
) -> "NeuralNetworkEnsemblePredictor":
"""Load a downloaded neural network ensemble artifact.
Args:
filename: Path to the zip artifact written by
:meth:`jcmoptimizer.ActiveLearning.download_neural_network`.
device: PyTorch device for prediction. If omitted, CUDA is used when
available and CPU otherwise.
Returns:
A predictor that evaluates the ensemble locally.
"""
with zipfile.ZipFile(filename, "r") as zf:
metadata = json.loads(zf.read("metadata.json").decode("utf-8"))
member_model_bytes = zf.read(metadata["member_model_file"])
member_state_bytes = zf.read(metadata["member_state_file"])
return cls(metadata, member_model_bytes, member_state_bytes, device=device)
def _infer_dtype(self, stacked_state: dict[str, torch.Tensor]) -> torch.dtype:
for value in stacked_state.values():
if self._torch.is_tensor(value) and value.is_floating_point():
return value.dtype
return self._torch.get_default_dtype()
[docs]
def predict(self, x: Union[torch.Tensor, list[list[float]]]) -> torch.Tensor:
"""Predict for one or more input points.
Args:
x: Input points with shape ``(num_points, num_input)``. For
one-dimensional input spaces, a one-dimensional sequence is also
accepted and interpreted as multiple points.
Returns:
Tensor with shape ``(num_models, num_points, num_output)``.
"""
x_tensor = self._torch.as_tensor(x, dtype=self.dtype, device=self.device)
if x_tensor.ndim == 1:
if self.num_input != 1:
raise ValueError(
f"Expected input with {self.num_input} columns, got a 1D input."
)
x_tensor = x_tensor.reshape(-1, 1)
if x_tensor.ndim != 2 or x_tensor.shape[1] != self.num_input:
raise ValueError(
f"Expected input of shape (num_points, {self.num_input}), "
f"got {tuple(x_tensor.shape)}."
)
return self._vmap(self.params, self.buffers, x_tensor)