From 771385da903d4410e78a0e7c9263b089e3f4f657 Mon Sep 17 00:00:00 2001
From: Yaman Umuroglu <maltanar@gmail.com>
Date: Wed, 16 Sep 2020 22:01:21 +0200
Subject: [PATCH] [Util] allow returning full context with execute_parent
---
src/finn/util/test.py | 8 +++++---
1 file changed, 5 insertions(+), 3 deletions(-)
diff --git a/src/finn/util/test.py b/src/finn/util/test.py
index 0b7284d76..32c6a0a3a 100644
--- a/src/finn/util/test.py
+++ b/src/finn/util/test.py
@@ -162,7 +162,7 @@ def get_trained_network_and_ishape(topology, wbits, abits):
return (model, ishape)
-def execute_parent(parent_path, child_path, input_tensor_npy):
+def execute_parent(parent_path, child_path, input_tensor_npy, return_full_ctx=False):
"""Execute parent model containing a single StreamingDataflowPartition by
replacing it with the model at child_path and return result."""
@@ -173,5 +173,7 @@ def execute_parent(parent_path, child_path, input_tensor_npy):
sdp_node = getCustomOp(sdp_node)
sdp_node.set_nodeattr("model", child_path)
ret = execute_onnx(parent_model, {iname: input_tensor_npy}, True)
- y = ret[oname]
- return y
+ if return_full_ctx:
+ return ret
+ else:
+ return ret[oname]
--
GitLab