ryanthompson591 commented on code in PR #22795:
URL: https://github.com/apache/beam/pull/22795#discussion_r956387116
##########
sdks/python/apache_beam/ml/inference/pytorch_inference.py:
##########
@@ -188,20 +216,24 @@ def __init__(
Otherwise, it will be CPU.
"""
self._state_dict_path = state_dict_path
- if device == 'GPU' and torch.cuda.is_available():
+ if device == 'GPU':
+ logging.info("Device is set to CUDA")
Review Comment:
I wonder if some of this logging is getting to be a little too much. Were
you using some of it to debug?
For example you have "device is set to" xxx in a couple spots. Is that still
important logging to keep?
##########
sdks/python/apache_beam/ml/inference/pytorch_inference.py:
##########
@@ -234,16 +266,17 @@ def run_inference(
# If elements in `batch` are provided as a dictionaries from key to
Tensors,
# then iterate through the batch list, and group Tensors to the same key
key_to_tensor_list = defaultdict(list)
- for example in batch:
- for key, tensor in example.items():
- key_to_tensor_list[key].append(tensor)
- key_to_batched_tensors = {}
- for key in key_to_tensor_list:
- batched_tensors = torch.stack(key_to_tensor_list[key])
- batched_tensors = _convert_to_device(batched_tensors, self._device)
- key_to_batched_tensors[key] = batched_tensors
- predictions = model(**key_to_batched_tensors, **inference_args)
- return [PredictionResult(x, y) for x, y in zip(batch, predictions)]
+ with torch.no_grad():
Review Comment:
since its not apparent why this is here. Consider linking the issue.
##########
sdks/python/apache_beam/ml/inference/pytorch_inference_test.py:
##########
@@ -373,6 +373,40 @@ def test_invalid_input_type(self):
# pylint: disable=expression-not-assigned
pcoll | RunInference(model_handler)
+ def test_gpu_convert_to_cpu(self):
+ with self.assertLogs() as log:
+ with TestPipeline() as pipeline:
+ examples = torch.from_numpy(
+ np.array([1, 5, 3, 10], dtype="float32").reshape(-1, 1))
+
+ state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
+ ('linear.bias', torch.Tensor([0.5]))])
+ path = os.path.join(self.tmpdir, 'my_state_dict_path')
+ torch.save(state_dict, path)
+
+ model_handler = PytorchModelHandlerTensor(
+ state_dict_path=path,
+ model_class=PytorchLinearRegression,
+ model_params={
+ 'input_dim': 1, 'output_dim': 1
+ },
+ device='GPU')
+ # Upon initialization, device is cuda
+ self.assertEqual(model_handler._device, torch.device('cuda'))
+
+ pcoll = pipeline | 'start' >> beam.Create(examples)
+ # pylint: disable=expression-not-assigned
+ pcoll | RunInference(model_handler)
+
+ # During model loading, device converted to cuda
+ self.assertEqual(model_handler._device, torch.device('cuda'))
+
+ self.assertIn("INFO:root:Device is set to CUDA", log.output)
+ self.assertIn(
Review Comment:
sure seems fine.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@beam.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
