class TestImplementations(object):
@pytest.mark.parametrize(
"obj",
[
Algorithm(
implementation=4,
tactic=5,
inputs=[(trt.TensorFormat.LINEAR, trt.float32)],
outputs=[(trt.TensorFormat.LINEAR, trt.float32)],
),
Algorithm(
implementation=4,
tactic=5,
inputs=[(trt.TensorFormat.LINEAR, trt.float32),
(trt.TensorFormat.CHW32, trt.int8)],
outputs=[(trt.TensorFormat.CHW32, trt.float16)],
),
np.ones((3, 4, 5), dtype=np.int64),
np.ones(5, dtype=np.int64),
np.zeros((4, 5), dtype=np.float32),
np.random.random_sample((3, 5)),
make_iter_result(),
RunResults([("runner0", [make_iter_result()]),
("runner0", [make_iter_result()])]),
],
ids=lambda x: type(x),
)
def test_serde(self, obj):
encoded = to_json(obj)
decoded = from_json(encoded)
if isinstance(obj, np.ndarray):
assert np.array_equal(decoded, obj)
else:
assert decoded == obj
@pytest.mark.parametrize("obj", JSONABLE_CASES)
def test_to_from_json(self, obj):
encoded = obj.to_json()
decoded = type(obj).from_json(encoded)
assert decoded == obj
@pytest.mark.parametrize("obj", JSONABLE_CASES)
def test_save_load(self, obj):
with tempfile.NamedTemporaryFile("w+") as f:
obj.save(f)
decoded = type(obj).load(f)
assert decoded == obj
def test_cannot_save_load_to_different_types(self):
run_result = JSONABLE_CASES[0]
encoded = run_result.to_json()
with pytest.raises(PolygraphyException,
match="JSON cannot be decoded into"):
TacticReplayData.from_json(encoded)
def replay(request):
"""
Returns:
Tuple[FakeAlgorithmContext, Algorithm, FakeAlgorithm,
Union[str, TacticReplayData], Union[str, TacticReplayData]]:
This fixture returns 5 things:
1. A fake TensorRT algorithm context
2. A Polygraphy Algorithm instance
3. A fake TensorRT algorithm (with the same information as (2))
4. An input tactic replay data, populated with the Polygraphy Algorithm (2), either
as a ``TacticReplayData`` instance, or a path.
5. An output tactic replay data, empty, either as a ``TacticReplayData`` instance, or
a path.
"""
jsonify = request.param
name = "node_of_y"
context = fake_context(name)
trt_algo = fake_algo()
poly_algo = Algorithm.from_trt(context, trt_algo)
in_replay_data = TacticReplayData().add(name, poly_algo)
out_replay_data = TacticReplayData()
if jsonify:
inpath = util.NamedTemporaryFile("w")
in_replay_data.save(inpath.name)
in_replay_data = inpath.name
outpath = util.NamedTemporaryFile("r")
out_replay_data = outpath.name
yield context, poly_algo, trt_algo, in_replay_data, out_replay_data
def make_algo():
return Algorithm(
implementation=4,
tactic=5,
inputs=[(trt.TensorFormat.LINEAR, trt.float32)],
outputs=[(trt.TensorFormat.LINEAR, trt.float32)],
)
# See the License for the specific language governing permissions and
# limitations under the License.
#
from collections import namedtuple
import pytest
import tensorrt as trt
from polygraphy import mod, util
from polygraphy.backend.trt import Algorithm, TacticRecorder, TacticReplayData, TacticReplayer
from polygraphy.exception import PolygraphyException
ALGO_EQ_CASES = [
(
Algorithm(6,
1,
inputs=[(trt.TensorFormat.LINEAR, trt.float32)],
outputs=[(trt.TensorFormat.LINEAR, trt.float32)]),
Algorithm(6,
1,
inputs=[(trt.TensorFormat.LINEAR, trt.float32)],
outputs=[(trt.TensorFormat.LINEAR, trt.float32)]),
True,
), # Same
(
Algorithm(6,
1,
inputs=[(trt.TensorFormat.LINEAR, trt.float32)],
outputs=[(trt.TensorFormat.LINEAR, trt.float32)]),
Algorithm(7,
1,
inputs=[(trt.TensorFormat.LINEAR, trt.float32)],
def make_replay(tactic):
return TacticReplayData().add(
"layer0",
Algorithm.from_trt(fake_context("layer0"), fake_algo(0, tactic)))