diff --git a/src/finn/util/test.py b/src/finn/util/test.py index 4cad01b1f7ec58da7ba6d5460c072faa01202c55..f29af66f52e18f8aeed87e19bfb52354ca1b73a7 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")