wehub-resource-sync
410 lines
14 KiB
Python
410 lines
14 KiB
Python
#
|
|
# 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.
|
|
#
|
|
|
|
from collections import OrderedDict
|
|
|
|
from polygraphy import config, mod, util
|
|
from polygraphy.common.interface import TypedDict, TypedList
|
|
from polygraphy.json import Decoder, Encoder, add_json_methods, load_json, save_json
|
|
from polygraphy.logger import G_LOGGER
|
|
|
|
|
|
class LazyArray:
|
|
"""
|
|
Represents a lazily loaded NumPy array or PyTorch Tensor.
|
|
For example, large arrays may be serialized to temporary files on the disk
|
|
to save memory.
|
|
"""
|
|
|
|
def __init__(self, arr):
|
|
"""
|
|
Args:
|
|
arr (Union[np.ndarray, torch.Tensor]): The array.
|
|
"""
|
|
self.arr = None
|
|
self.tmpfile = None
|
|
if config.ARRAY_SWAP_THRESHOLD_MB >= 0 and util.array.nbytes(arr) > (
|
|
config.ARRAY_SWAP_THRESHOLD_MB << 20
|
|
):
|
|
self.tmpfile = util.NamedTemporaryFile(suffix=".json")
|
|
G_LOGGER.extra_verbose(
|
|
f"Evicting large array ({util.array.nbytes(arr) / 1024.0 ** 2:.3f} MiB) from memory and saving to {self.tmpfile.name}"
|
|
)
|
|
save_json(arr, self.tmpfile.name)
|
|
else:
|
|
self.arr = arr
|
|
|
|
def load(self):
|
|
"""
|
|
Load the array, deserializing from the disk if it was stored earlier.
|
|
|
|
Returns:
|
|
Union[np.ndarray, torch.Tensor]: The array
|
|
"""
|
|
if self.arr is not None:
|
|
return self.arr
|
|
|
|
if self.tmpfile is None:
|
|
G_LOGGER.internal_error(
|
|
f"self.arr is None but self.tmpfile is also None; this should be impossible."
|
|
)
|
|
return load_json(self.tmpfile.name)
|
|
|
|
|
|
@Encoder.register(LazyArray, alias="LazyNumpyArray")
|
|
def encode(lazy_arr):
|
|
return {
|
|
"values": lazy_arr.load(),
|
|
}
|
|
|
|
|
|
@Decoder.register(LazyArray, alias="LazyNumpyArray")
|
|
def decode(dct):
|
|
return LazyArray(dct["values"])
|
|
|
|
|
|
@mod.export()
|
|
class IterationResult(TypedDict(lambda: str, lambda: LazyArray)):
|
|
"""
|
|
An ordered dictionary containing the result of a running a single iteration of a runner.
|
|
|
|
This maps output names to arrays, and preserves the output ordering from the runner.
|
|
|
|
NOTE: The ``POLYGRAPHY_ARRAY_SWAP_THRESHOLD_MB`` environment variable can be set to enable
|
|
the arrays to be swapped to the disk.
|
|
|
|
Also includes additional fields indicating the name of the runner which produced the
|
|
outputs, and the time required to do so.
|
|
"""
|
|
|
|
@staticmethod
|
|
def _to_lazy(nparray):
|
|
if isinstance(nparray, LazyArray):
|
|
return nparray
|
|
return LazyArray(nparray)
|
|
|
|
@staticmethod
|
|
def _to_lazy_dict(nparray_dict):
|
|
if nparray_dict is None:
|
|
return None
|
|
|
|
# Converts a Dict[str, np.ndarray] to a Dict[str, LazyArray]
|
|
lazy = OrderedDict()
|
|
for name, out in nparray_dict.items():
|
|
lazy[name] = IterationResult._to_lazy(out)
|
|
return lazy
|
|
|
|
def __init__(self, outputs=None, runtime=None, runner_name=None):
|
|
"""
|
|
Args:
|
|
outputs (Dict[str, Union[np.array, torch.Tensor]]): The outputs of this iteration, mapped to their names.
|
|
|
|
runtime (float):
|
|
The time required for this iteration, in seconds.
|
|
Only used for logging purposes.
|
|
runner_name (str):
|
|
The name of the runner that produced this output.
|
|
If this is omitted, a default name is generated.
|
|
"""
|
|
if outputs and config.ARRAY_SWAP_THRESHOLD_MB < 0:
|
|
total_size_gb = sum(
|
|
util.array.nbytes(arr)
|
|
for arr in outputs.values()
|
|
if util.array.is_torch(arr) or util.array.is_numpy(arr)
|
|
) / (1024.0**3)
|
|
if total_size_gb >= 1:
|
|
G_LOGGER.warning(
|
|
f"It looks like the outputs of this network are very large ({total_size_gb:.3f} GiB).\n"
|
|
"To reduce memory usage, you may want to allow Polygraphy to swap these arrays to the disk "
|
|
"using the POLYGRAPHY_ARRAY_SWAP_THRESHOLD_MB environment variable."
|
|
)
|
|
|
|
super().__init__(IterationResult._to_lazy_dict(outputs))
|
|
self.runtime = runtime
|
|
self.runner_name = util.default(runner_name, "custom_runner")
|
|
|
|
# Convenience methods to preserve np.ndarray in the interface.
|
|
def update(self, other):
|
|
return super().update(IterationResult._to_lazy_dict(other))
|
|
|
|
def __setitem__(self, name, arr):
|
|
return super().__setitem__(name, IterationResult._to_lazy(arr))
|
|
|
|
def values(self):
|
|
for arr in super().values():
|
|
yield arr.load()
|
|
|
|
def items(self):
|
|
for name, arr in super().items():
|
|
yield name, arr.load()
|
|
|
|
def __getitem__(self, name):
|
|
return super().__getitem__(name).load()
|
|
|
|
def __eq__(self, other):
|
|
if self.runtime != other.runtime or self.runner_name != other.runner_name:
|
|
return False
|
|
|
|
for key, val in self.items():
|
|
if key not in other:
|
|
return False
|
|
|
|
if not util.array.equal(val, other[key]):
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
@Encoder.register(IterationResult)
|
|
def encode(iter_result):
|
|
return {
|
|
"outputs": iter_result.dct,
|
|
"runtime": iter_result.runtime,
|
|
"runner_name": iter_result.runner_name,
|
|
}
|
|
|
|
|
|
@Decoder.register(IterationResult)
|
|
def decode(dct):
|
|
return IterationResult(
|
|
outputs=dct["outputs"], runtime=dct["runtime"], runner_name=dct["runner_name"]
|
|
)
|
|
|
|
|
|
@mod.export()
|
|
@add_json_methods("inference results")
|
|
class RunResults(TypedList(lambda: tuple)):
|
|
"""
|
|
Maps runners to per-iteration outputs (in the form of a ``List[IterationResult]``).
|
|
|
|
For example, if ``results`` is an instance of ``RunResults()``, then
|
|
to access the outputs of the first iteration from a specified runner, do:
|
|
::
|
|
|
|
iteration = 0
|
|
runner_name = "trt-runner"
|
|
outputs = results[runner_name][iteration]
|
|
|
|
# `outputs` is a `Dict[str, np.ndarray]`
|
|
|
|
|
|
Note: Technically, this is a ``List[Tuple[str, List[IterationResult]]]``, but includes
|
|
helpers that make it behave like an OrderedDict that can contain duplicates.
|
|
"""
|
|
|
|
def items(self):
|
|
"""
|
|
Creates a generator that yields ``Tuple[str, List[IterationResult]]`` - runner names
|
|
and corresponding outputs.
|
|
"""
|
|
for name, iteration_results in self.lst:
|
|
yield name, iteration_results
|
|
|
|
def keys(self):
|
|
"""
|
|
Creates a generator that yields runner names (str).
|
|
"""
|
|
for name, _ in self.lst:
|
|
yield name
|
|
|
|
def values(self):
|
|
"""
|
|
Creates a generator that yields runner outputs (List[IterationResult]).
|
|
"""
|
|
for _, iteration_results in self.lst:
|
|
yield iteration_results
|
|
|
|
def update(self, other):
|
|
"""
|
|
Updates the results stored in this instance.
|
|
|
|
Args:
|
|
other (Union[Dict[str, List[IterationResult]], RunResults]):
|
|
A dictionary or RunResults instance from which to update this one.
|
|
"""
|
|
for name, iteration_results in other.items():
|
|
self.lst[name] = iteration_results
|
|
return self
|
|
|
|
def add(self, out_list, runtime=None, runner_name=None):
|
|
"""
|
|
A helper to create a ``List[IterationResult]`` and map it to the specified runner_name.
|
|
|
|
This method cannot be used to modify an existing entry.
|
|
|
|
Calling this method is equivalent to:
|
|
::
|
|
|
|
results[runner_name] = []
|
|
for out in out_list:
|
|
results[runner_name].append(IterationResult(out, runtime, runner_name))
|
|
|
|
Args:
|
|
out_list (List[Dict[str, np.array]]):
|
|
One or more set of outputs where each output is a dictionary
|
|
of output names mapped to NumPy arrays.
|
|
|
|
runtime (float):
|
|
The time required for this iteration, in seconds.
|
|
Only used for logging purposes.
|
|
runner_name (str):
|
|
The name of the runner that produced this output.
|
|
If this is omitted, a default name is generated.
|
|
"""
|
|
runner_name = util.default(runner_name, "custom_runner")
|
|
iter_results = [IterationResult(out, runtime, runner_name) for out in out_list]
|
|
self[runner_name] = iter_results
|
|
|
|
def __getitem__(self, key):
|
|
if isinstance(key, int):
|
|
return self.lst[key]
|
|
|
|
for name, iteration_results in self.lst:
|
|
if name == key:
|
|
return iteration_results
|
|
|
|
G_LOGGER.critical(
|
|
f"{key:35} does not exist in this RunResults instance. Note: Available runners: {list(self.keys())}"
|
|
)
|
|
|
|
def __setitem__(self, key, value):
|
|
if isinstance(key, int):
|
|
self.lst[key] = value
|
|
return
|
|
|
|
for index, name in enumerate(self.keys()):
|
|
if name == key:
|
|
self.lst[index] = (key, value)
|
|
break
|
|
else:
|
|
self.append((key, value))
|
|
|
|
def __contains__(self, val):
|
|
if isinstance(val, str) or isinstance(val, bytes):
|
|
return val in list(self.keys())
|
|
return val in self.lst
|
|
|
|
def __eq__(self, other):
|
|
for (r0, its0), (r1, its1) in zip(self.lst, other.lst):
|
|
if r0 != r1:
|
|
return False
|
|
|
|
if its0 != its1:
|
|
return False
|
|
return True
|
|
|
|
|
|
@Encoder.register(RunResults)
|
|
def encode(results):
|
|
return {"lst": results.lst}
|
|
|
|
|
|
@Decoder.register(RunResults)
|
|
def decode(dct):
|
|
return RunResults(list(map(tuple, dct["lst"])))
|
|
|
|
|
|
@mod.export()
|
|
class AccuracyResult(TypedDict(lambda: tuple, lambda: list)):
|
|
"""
|
|
An ordered dictionary including details about the result of ``Comparator.compare_accuracy``.
|
|
|
|
More specifically, it is an ``OrderedDict[Tuple[str, str], List[OrderedDict[str, bool]]]`` which maps a runner
|
|
pair (a tuple containing both runner names) to a list of dictionaries of booleans (or anything that can be
|
|
converted into a boolean, such as an ``OutputCompareResult``), indicating whether there was a match in the outputs of
|
|
the corresponding iteration. The ``List[OrderedDict[str, bool]]`` is constructed from the dictionaries returned
|
|
by ``compare_func`` in ``compare_accuracy``.
|
|
|
|
For example, to see if there's a match between ``runner0`` and
|
|
``runner1`` during the 1st iteration for an output called ``output0``:
|
|
::
|
|
|
|
runner_pair = ("runner0", "runner1")
|
|
iteration = 0
|
|
output_name = "output0"
|
|
match = bool(accuracy_result[runner_pair][iteration][output_name])
|
|
|
|
If there's a mismatch, you can inspect the outputs from
|
|
the results of ``Comparator.run()``, assumed here to be called ``run_results``:
|
|
::
|
|
|
|
runner0_output = run_results["runner0"][iteration][output_name]
|
|
runner1_output = run_results["runner1"][iteration][output_name]
|
|
"""
|
|
|
|
def __bool__(self):
|
|
"""
|
|
Whether all outputs matched for every iteration.
|
|
You can use this function to avoid manually checking each output. For example:
|
|
::
|
|
|
|
if accuracy_result:
|
|
print("All matched!")
|
|
|
|
Returns:
|
|
bool
|
|
"""
|
|
return all(
|
|
[
|
|
bool(match)
|
|
for outs in self.values()
|
|
for out in outs
|
|
for match in out.values()
|
|
]
|
|
)
|
|
|
|
def _get_runner_pair(self, runner_pair):
|
|
return util.default(runner_pair, list(self.keys())[0])
|
|
|
|
def percentage(self, runner_pair=None):
|
|
"""
|
|
Returns the percentage of iterations that matched for the given pair of runners,
|
|
expressed as a decimal between 0.0 and 1.0.
|
|
|
|
Always returns 1.0 when the number of iterations is 0, or when there are no runner comparisons.
|
|
|
|
Args:
|
|
runner_pair (Tuple[str, str]):
|
|
A pair of runner names describing which runners to check.
|
|
Defaults to the first pair in the dictionary.
|
|
"""
|
|
if not list(self.keys()):
|
|
return 1.0 # No data in this result.
|
|
|
|
matched, _, total = self.stats(runner_pair)
|
|
if not total:
|
|
return 1.0 # No iterations
|
|
return float(matched) / float(total)
|
|
|
|
def stats(self, runner_pair=None):
|
|
"""
|
|
Returns the number of iterations that matched, mismatched, and the total number of iterations.
|
|
|
|
Args:
|
|
runner_pair (Tuple[str, str]):
|
|
A pair of runner names describing which runners to check.
|
|
Defaults to the first pair in the dictionary.
|
|
|
|
Returns:
|
|
Tuple[int, int, int]: Number of iterations that matched, mismatched, and total respectively.
|
|
"""
|
|
runner_pair = self._get_runner_pair(runner_pair)
|
|
outs = self[runner_pair]
|
|
matched = sum([all([match for match in out.values()]) for out in outs])
|
|
total = len(outs)
|
|
return matched, total - matched, total
|