chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,258 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
import copy
|
||||
import time
|
||||
from collections import defaultdict
|
||||
|
||||
from polygraphy import config, func, mod, util
|
||||
from polygraphy.datatype import DataType
|
||||
from polygraphy.logger import G_LOGGER, LogMode
|
||||
from polygraphy.backend.base import util as base_util
|
||||
|
||||
np = mod.lazy_import("numpy")
|
||||
|
||||
|
||||
@mod.export()
|
||||
class BaseRunner:
|
||||
"""
|
||||
Base class for Polygraphy runners. All runners should override the functions and attributes specified here.
|
||||
"""
|
||||
|
||||
RUNNER_COUNTS = defaultdict(int)
|
||||
|
||||
def __init__(self, name=None, prefix=None):
|
||||
"""
|
||||
Args:
|
||||
name (str):
|
||||
The name to use for this runner.
|
||||
prefix (str):
|
||||
The human-readable name prefix to use for this runner.
|
||||
A runner count and timestamp will be appended to this prefix.
|
||||
Only used if name is not provided.
|
||||
"""
|
||||
prefix = util.default(prefix, "Runner")
|
||||
if name is None:
|
||||
count = BaseRunner.RUNNER_COUNTS[prefix]
|
||||
BaseRunner.RUNNER_COUNTS[prefix] += 1
|
||||
name = f"{prefix}-N{count}-{time.strftime('%x')}-{time.strftime('%X')}"
|
||||
self.name = name
|
||||
self.inference_time = None
|
||||
|
||||
self.is_active = False
|
||||
"""bool: Whether this runner has been activated, either via context manager, or by calling ``activate()``."""
|
||||
|
||||
def __enter__(self):
|
||||
"""
|
||||
Activate the runner for inference. For example, this may involve allocating CPU or GPU memory.
|
||||
"""
|
||||
self.activate()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
"""
|
||||
Deactivate the runner. For example, this may involve freeing CPU or GPU memory.
|
||||
"""
|
||||
self.deactivate()
|
||||
|
||||
# Implementation for runner activation. Derived classes should override this function
|
||||
# rather than ``activate()``.
|
||||
def activate_impl(self):
|
||||
pass
|
||||
|
||||
def activate(self):
|
||||
"""
|
||||
Activate the runner for inference. For example, this may involve allocating CPU or GPU memory.
|
||||
|
||||
Generally, you should use a context manager instead of manually activating and deactivating.
|
||||
For example:
|
||||
::
|
||||
|
||||
with RunnerType(...) as runner:
|
||||
runner.infer(...)
|
||||
"""
|
||||
if self.is_active:
|
||||
G_LOGGER.warning(
|
||||
f"{self.name:35} | Already active; will not activate again. "
|
||||
"If you really want to activate this runner again, call activate_impl() directly"
|
||||
)
|
||||
return
|
||||
|
||||
if config.INTERNAL_CORRECTNESS_CHECKS:
|
||||
self._pre_activate_runner_state = copy.copy(vars(self))
|
||||
|
||||
self.activate_impl()
|
||||
self.is_active = True
|
||||
|
||||
def get_input_metadata_impl(self):
|
||||
"""
|
||||
Implemenation for `get_input_metadata`. Derived classes should override this function
|
||||
rather than `get_input_metadata`.
|
||||
|
||||
Derived classes may return any kind of data type supported by Polygraphy's DataType
|
||||
class (e.g. np.dtype, torch.dtype, etc.)
|
||||
"""
|
||||
raise NotImplementedError("BaseRunner is an abstract class")
|
||||
|
||||
@func.constantmethod
|
||||
def get_input_metadata(self, use_numpy_dtypes=None):
|
||||
"""
|
||||
Returns information about the inputs of the model.
|
||||
Shapes here may include dynamic dimensions, represented by ``None``.
|
||||
Must be called only after ``activate()`` and before ``deactivate()``.
|
||||
|
||||
Args:
|
||||
use_numpy_dtypes (bool):
|
||||
[DEPRECATED] Whether to return NumPy data types instead of Polygraphy ``DataType`` s.
|
||||
This is provided to retain backwards compatibility. In the future,
|
||||
this parameter will be removed and Polygraphy ``DataType`` s will
|
||||
always be returned. These can be converted to NumPy data types by calling the `numpy()` method.
|
||||
Defaults to True.
|
||||
|
||||
Returns:
|
||||
TensorMetadata: Input names, shapes, and data types.
|
||||
"""
|
||||
if not self.is_active:
|
||||
G_LOGGER.critical(
|
||||
f"{self.name:35} | Must be activated prior to calling get_input_metadata()"
|
||||
)
|
||||
|
||||
use_numpy_dtypes = util.default(use_numpy_dtypes, True)
|
||||
|
||||
meta = self.get_input_metadata_impl()
|
||||
|
||||
for name, (dtype, _) in meta.items():
|
||||
dtype = DataType.from_dtype(dtype)
|
||||
if use_numpy_dtypes:
|
||||
mod.warn_deprecated(
|
||||
"Returning NumPy data types instead of Polygraphy `DataType`s from `get_input_metadata()`",
|
||||
use_instead=None,
|
||||
remove_in="0.60.0",
|
||||
)
|
||||
meta[name]._dtype = DataType.to_dtype(dtype, "numpy")
|
||||
return meta
|
||||
|
||||
# Implementation for runner inference. Derived classes should override this function
|
||||
# rather than ``infer()``
|
||||
# Derived classes should also set the `inference_time` property so that performance metrics are accurate.
|
||||
def infer_impl(self, feed_dict):
|
||||
raise NotImplementedError("BaseRunner is an abstract class")
|
||||
|
||||
def infer(self, feed_dict, check_inputs=True, *args, **kwargs):
|
||||
"""
|
||||
Runs inference using the provided feed_dict.
|
||||
|
||||
Must be called only after ``activate()`` and before ``deactivate()``.
|
||||
|
||||
NOTE: Some runners may accept additional parameters in infer().
|
||||
For details on these, see the documentation for their `infer_impl()` methods.
|
||||
|
||||
Args:
|
||||
feed_dict (OrderedDict[str, numpy.ndarray]):
|
||||
A mapping of input tensor names to corresponding input NumPy arrays.
|
||||
|
||||
check_inputs (bool):
|
||||
Whether to check that the provided ``feed_dict`` includes the expected inputs
|
||||
with the expected data types and shapes.
|
||||
Disabling this may improve performance.
|
||||
Defaults to True.
|
||||
|
||||
Attributes:
|
||||
inference_time (float):
|
||||
The time required to run inference in seconds.
|
||||
|
||||
Returns:
|
||||
OrderedDict[str, numpy.ndarray]:
|
||||
A mapping of output tensor names to their corresponding NumPy arrays.
|
||||
|
||||
IMPORTANT: Runners may reuse these output buffers. Thus, if you need to save
|
||||
outputs from multiple inferences, you should make a copy with ``copy.deepcopy(outputs)``.
|
||||
"""
|
||||
if not self.is_active:
|
||||
G_LOGGER.critical(
|
||||
f"{self.name:35} | Must be activated prior to calling infer()"
|
||||
)
|
||||
|
||||
if check_inputs:
|
||||
input_metadata = self.get_input_metadata(use_numpy_dtypes=False)
|
||||
G_LOGGER.verbose(
|
||||
f"{self.name:35} | Input metadata is: {input_metadata}",
|
||||
mode=LogMode.ONCE,
|
||||
)
|
||||
base_util.check_inputs(feed_dict, input_metadata)
|
||||
|
||||
return self.infer_impl(feed_dict, *args, **kwargs)
|
||||
|
||||
@func.constantmethod
|
||||
def last_inference_time(self):
|
||||
"""
|
||||
Returns the total inference time in seconds required during the last call to ``infer()``.
|
||||
|
||||
Must be called only after ``activate()`` and before ``deactivate()``.
|
||||
|
||||
Returns:
|
||||
float: The time in seconds, or None if runtime was not measured by the runner.
|
||||
"""
|
||||
if self.inference_time is None:
|
||||
msg = f"{self.name:35} | `inference_time` was not set. Inference time will be incorrect! "
|
||||
msg += "To correctly compare runtimes, please set the `inference_time` attribute in `infer_impl()`"
|
||||
|
||||
G_LOGGER.internal_error(msg)
|
||||
G_LOGGER.warning(msg, mode=LogMode.ONCE)
|
||||
return None
|
||||
return self.inference_time
|
||||
|
||||
# Implementation for runner deactivation. Derived classes should override this function
|
||||
# rather than ``deactivate()``.
|
||||
def deactivate_impl(self):
|
||||
pass
|
||||
|
||||
def deactivate(self):
|
||||
"""
|
||||
Deactivate the runner. For example, this may involve freeing CPU or GPU memory.
|
||||
|
||||
Generally, you should use a context manager instead of manually activating and deactivating.
|
||||
For example:
|
||||
::
|
||||
|
||||
with RunnerType(...) as runner:
|
||||
runner.infer(...)
|
||||
"""
|
||||
if not self.is_active:
|
||||
G_LOGGER.warning(
|
||||
f"{self.name:35} | Not active; will not deactivate. If you really want to deactivate this runner, call deactivate_impl() directly"
|
||||
)
|
||||
return
|
||||
|
||||
self.inference_time = None
|
||||
self.is_active = None
|
||||
|
||||
self.deactivate_impl()
|
||||
self.is_active = False
|
||||
if config.INTERNAL_CORRECTNESS_CHECKS:
|
||||
old_state = self._pre_activate_runner_state
|
||||
del self._pre_activate_runner_state
|
||||
if old_state != vars(self):
|
||||
G_LOGGER.internal_error(
|
||||
f"Runner state was not reset after deactivation. Note:\nOld state: {old_state}\nNew state: {vars(self)}"
|
||||
)
|
||||
|
||||
def __del__(self):
|
||||
if self.is_active:
|
||||
# __del__ is not guaranteed to be called, but when it is, this could be a useful warning.
|
||||
print(
|
||||
f"[W] {self.name:35} | Was activated but never deactivated. This could cause a memory leak!"
|
||||
)
|
||||
Reference in New Issue
Block a user