From e2507c37152d3e3318528c192a83737be3c0e656 Mon Sep 17 00:00:00 2001
From: Yaman Umuroglu <maltanar@gmail.com>
Date: Thu, 14 May 2020 17:32:03 +0100
Subject: [PATCH] [Util] introduce a load_test_checkpoint_or_skip fxn
---
src/finn/util/test.py | 13 +++++++++++++
1 file changed, 13 insertions(+)
diff --git a/src/finn/util/test.py b/src/finn/util/test.py
index 4cad01b1f..f29af66f5 100644
--- a/src/finn/util/test.py
+++ b/src/finn/util/test.py
@@ -27,6 +27,9 @@
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
from brevitas_examples import bnn_pynq
+import pytest
+import warnings
+from finn.core.modelwrapper import ModelWrapper
# map of (wbits,abits) -> model
example_map = {
@@ -61,3 +64,13 @@ def get_test_model_trained(netname, wbits, abits):
def get_test_model_untrained(netname, wbits, abits):
"get_test_model with pretrained=False"
return get_test_model(netname, wbits, abits, pretrained=False)
+
+
+def load_test_checkpoint_or_skip(filename):
+ "Try to load given .onnx and return ModelWrapper, else skip current test."
+ try:
+ model = ModelWrapper(filename)
+ return model
+ except FileNotFoundError:
+ warnings.warn(filename + " not found from previous test step, skipping")
+ pytest.skip(filename + " not found from previous test step, skipping")
--
GitLab