def test_pipeline_local_model(self):
with TestPipeline() as pipeline:
examples = torch.from_numpy(
np.array([1, 5, 3, 10, -14, 0, 0.5, 0.5],
dtype="float32")).reshape(-1, 2)
expected_predictions = [
PredictionResult(ex, pred) for ex, pred in zip(
examples,
torch.Tensor(
[f1 * 2.0 + f2 * 3 + 0.5
for f1, f2 in examples]).reshape(-1, 1))
]
state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0,
3]])),
('linear.bias', torch.Tensor([0.5]))])
path = os.path.join(self.tmpdir, 'my_state_dict_path')
torch.save(state_dict, path)
model_loader = PytorchModelLoader(
state_dict_path=path,
model_class=PytorchLinearRegression(input_dim=2, output_dim=1))
pcoll = pipeline | 'start' >> beam.Create(examples)
predictions = pcoll | RunInference(model_loader)
assert_that(
predictions,
equal_to(expected_predictions,
equals_fn=_compare_prediction_result))
def test_inference_runner_multiple_tensor_features(self):
examples = torch.from_numpy(
np.array([1, 5, 3, 10, -14, 0, 0.5, 0.5],
dtype="float32")).reshape(-1, 2)
examples = [
torch.from_numpy(np.array([1, 5], dtype="float32")),
torch.from_numpy(np.array([3, 10], dtype="float32")),
torch.from_numpy(np.array([-14, 0], dtype="float32")),
torch.from_numpy(np.array([0.5, 0.5], dtype="float32")),
]
expected_predictions = [
PredictionResult(ex, pred) for ex, pred in zip(
examples,
torch.Tensor([f1 * 2.0 + f2 * 3 + 0.5
for f1, f2 in examples]).reshape(-1, 1))
]
model = PytorchLinearRegression(input_dim=2, output_dim=1)
model.load_state_dict(
OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])),
('linear.bias', torch.Tensor([0.5]))]))
model.eval()
inference_runner = PytorchInferenceRunner(torch.device('cpu'))
predictions = inference_runner.run_inference(examples, model)
for actual, expected in zip(predictions, expected_predictions):
self.assertTrue(_compare_prediction_result(actual, expected))
def test_pipeline_gcs_model(self):
with TestPipeline() as pipeline:
examples = torch.from_numpy(
np.array([1, 5, 3, 10], dtype="float32").reshape(-1, 1))
expected_predictions = [
PredictionResult(ex, pred) for ex, pred in zip(
examples,
torch.Tensor([example * 2.0 + 0.5
for example in examples]).reshape(-1, 1))
]
gs_pth = 'gs://apache-beam-ml/pytorch_lin_reg_model_2x+0.5_state_dict.pth'
model_loader = PytorchModelLoader(
state_dict_path=gs_pth,
model_class=PytorchLinearRegression,
model_params={
'input_dim': 1,
'output_dim': 1
})
pcoll = pipeline | 'start' >> beam.Create(examples)
predictions = pcoll | RunInference(model_loader)
assert_that(
predictions,
equal_to(expected_predictions,
equals_fn=_compare_prediction_result))
def run_inference(self, batch: List[torch.Tensor],
model: torch.nn.Module) -> Iterable[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.
"""
batch = torch.stack(batch)
if batch.device != self._device:
batch = batch.to(self._device)
predictions = model(batch)
return [PredictionResult(x, y) for x, y in zip(batch, predictions)]
def run_inference(self, batch: List[numpy.ndarray],
model: Any) -> Iterable[PredictionResult]:
# vectorize data for better performance
vectorized_batch = numpy.stack(batch, axis=0)
predictions = model.predict(vectorized_batch)
return [PredictionResult(x, y) for x, y in zip(batch, predictions)]