apache_beam.ml.inference.pytorch module

class apache_beam.ml.inference.pytorch.PytorchInferenceRunner(device: torch.device)

Bases: apache_beam.ml.inference.base.InferenceRunner

This class runs Pytorch inferences with the run_inference method. It also has other methods to get the bytes of a batch of Tensors as well as the namespace for Pytorch models.

run_inference(batch: List[torch.Tensor], model: torch.nn.modules.module.Module) → Iterable[apache_beam.ml.inference.api.PredictionResult]

Runs inferences on a batch of Tensors and returns an Iterable of Tensor Predictions.

This method stacks the list of Tensors in a vectorized format to optimize the inference call.

get_num_bytes(batch: List[torch.Tensor]) → int

Returns the number of bytes of data for a batch of Tensors.

get_metrics_namespace() → str

Returns a namespace for metrics collected by the RunInference transform.

class apache_beam.ml.inference.pytorch.PytorchModelLoader(state_dict_path: str, model_class: Callable[[...], torch.nn.modules.module.Module], model_params: Dict[str, Any], device: str = 'CPU')

Bases: apache_beam.ml.inference.base.ModelLoader

Loads a Pytorch Model.

Initializes a PytorchModelLoader :param state_dict_path: path to the saved dictionary of the model state. :param model_class: class of the Pytorch model that defines the model structure. :param device: the device on which you wish to run the model. If device = GPU then a GPU device will be used if it is available. Otherwise, it will be CPU.

See https://pytorch.org/tutorials/beginner/saving_loading_models.html for details

load_model() → torch.nn.modules.module.Module

Loads and initializes a Pytorch model for processing.

get_inference_runner() → apache_beam.ml.inference.base.InferenceRunner

Returns a Pytorch implementation of InferenceRunner.