def test_custom_outputs(self):
loader = ModifyOnnx(OnnxFromPath(
ONNX_MODELS["identity_identity"].path),
outputs=["identity_out_0"])
model = loader()
assert len(model.graph.output) == 1
assert model.graph.output[0].name == "identity_out_0"
def test_save_onnx(self):
with tempfile.TemporaryDirectory() as outdir:
outpath = os.path.join(outdir, "test", "nested")
loader = SaveOnnx(OnnxFromPath(ONNX_MODELS["identity"].path),
path=outpath)
loader()
assert is_file_non_empty(outpath)
def test_exclude_outputs_with_layerwise(self):
loader = ModifyOnnx(OnnxFromPath(
ONNX_MODELS["identity_identity"].path),
outputs=constants.MARK_ALL,
exclude_outputs=["identity_out_2"])
model = loader()
assert len(model.graph.output) == 1
assert model.graph.output[0].name == "identity_out_0"
def test_external_data(self):
with util.NamedTemporaryFile() as path, util.NamedTemporaryFile(
) as data:
model = OnnxFromPath(ONNX_MODELS["const_foldable"].path)
loader = SaveOnnx(model,
path.name,
external_data_path=data.name,
size_threshold=0)
loader()
assert is_file_non_empty(path.name)
assert is_file_non_empty(data.name)
def test_basic(self):
model = ONNX_MODELS["identity"]
with OnnxTfRunner(OnnxFromPath(model.path)) as runner:
assert runner.is_active
model.check_runner(runner)
assert not runner.is_active
def test_save_onnx(self):
with tempfile.NamedTemporaryFile() as outpath:
loader = SaveOnnx(OnnxFromPath(ONNX_MODELS["identity"].path),
path=outpath.name)
loader()
check_file_non_empty(outpath.name)
def test_layerwise(self):
loader = ModifyOnnx(OnnxFromPath(
ONNX_MODELS["identity_identity"].path),
outputs=constants.MARK_ALL)
model = loader()
assert len(model.graph.output) == 2
def test_basic(self):
loader = OnnxFromPath(ONNX_MODELS["identity"].path)
loader()
def test_basic(self):
graph = gs_from_onnx(OnnxFromPath(ONNX_MODELS["identity"].path))
assert isinstance(graph, gs.Graph)
def test_external_data(self):
model = ONNX_MODELS["ext_weights"]
loader = OnnxFromPath(model.path, model.ext_data)
assert isinstance(loader(), onnx.ModelProto)
def test_basic(self):
loader = OnnxFromPath(ONNX_MODELS["identity"].path)
assert isinstance(loader(), onnx.ModelProto)