pranavm-nvidia commented on code in PR #22131:
URL: https://github.com/apache/beam/pull/22131#discussion_r922182382
##########
sdks/python/apache_beam/ml/inference/tensorrt_inference.py:
##########
@@ -80,9 +79,70 @@ def _validate_inference_args(inference_args):
'engines do not need extra arguments in their execute_v2() call.')
+class TensorRTEngine:
+ def __init__(self, engine: trt.ICudaEngine):
+ """Implementation of the TensorRTEngine class which handles
+ allocations associated with TensorRT engine.
+
+ Example Usage::
+
+ TensorRTEngine(engine)
+
+ Args:
+ engine: trt.ICudaEngine object that contains TensorRT engine
+ """
+ self.engine = engine
+ self.context = engine.create_execution_context()
+ self.inputs = []
+ self.outputs = []
+ self.gpu_allocations = []
+ self.cpu_allocations = []
+
+ # Setup I/O bindings
+ for i in range(self.engine.num_bindings):
+ is_input = False
+ if self.engine.binding_is_input(i):
+ is_input = True
+ name = self.engine.get_binding_name(i)
+ dtype = self.engine.get_binding_dtype(i)
+ shape = self.engine.get_binding_shape(i)
+ if is_input:
+ batch_size = shape[0]
+ size = np.dtype(trt.nptype(dtype)).itemsize
+ for s in shape:
+ size *= s
+ err, allocation = cuda.cuMemAlloc(size)
+ binding = {
+ 'index': i,
+ 'name': name,
+ 'dtype': np.dtype(trt.nptype(dtype)),
+ 'shape': list(shape),
+ 'allocation': allocation,
+ 'size': size
+ }
+ self.gpu_allocations.append(allocation)
+ if self.engine.binding_is_input(i):
+ self.inputs.append(binding)
+ else:
+ self.outputs.append(binding)
+
+ assert self.context
+ assert batch_size > 0
Review Comment:
Is `batch_size` used anywhere? I think it can be removed otherwise
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
