diff --git a/src/finn/transformation/fpgadataflow/cleanup.py b/src/finn/transformation/fpgadataflow/cleanup.py new file mode 100644 index 0000000000000000000000000000000000000000..1632d3443a3bf79e55a4b877ae182964ff7caaed --- /dev/null +++ b/src/finn/transformation/fpgadataflow/cleanup.py @@ -0,0 +1,35 @@ +import os +import shutil + +import finn.core.utils as util +import finn.custom_op.registry as registry +from finn.transformation import Transformation + + +class CleanUp(Transformation): + """Remove any generated files for fpgadataflow nodes.""" + + def __init__(self): + super().__init__() + + def apply(self, model): + for node in model.graph.node: + op_type = node.op_type + if node.domain == "finn": + backend_attribute = util.get_by_name(node.attribute, "backend") + backend_value = backend_attribute.s.decode("UTF-8") + if backend_value == "fpgadataflow": + try: + # lookup op_type in registry of CustomOps + inst = registry.custom_op[op_type](node) + code_gen_dir = inst.get_nodeattr("code_gen_dir") + if os.path.isdir(code_gen_dir): + shutil.rmtree(code_gen_dir) + inst.set_nodeattr("code_gen_dir", "") + inst.set_nodeattr("executable_path", "") + except KeyError: + # exception if op_type is not supported + raise Exception( + "Custom op_type %s is currently not supported." % op_type + ) + return (model, False) diff --git a/tests/fpgadataflow/test_code_gen_trafo.py b/tests/fpgadataflow/test_code_gen_trafo.py index e2651acfd0345156809e9c82f55694c0d48af4f9..53f74ffcbd72cb1338bc39cbe012880b892fc737 100644 --- a/tests/fpgadataflow/test_code_gen_trafo.py +++ b/tests/fpgadataflow/test_code_gen_trafo.py @@ -5,6 +5,7 @@ from onnx import TensorProto, helper import finn.core.utils as util from finn.core.datatype import DataType from finn.core.modelwrapper import ModelWrapper +from finn.transformation.fpgadataflow.cleanup import CleanUp from finn.transformation.fpgadataflow.codegen import CodeGen @@ -69,3 +70,4 @@ def test_code_gen_trafo(): op type {} is empty!""".format( node.op_type ) + model = model.transform(CleanUp()) diff --git a/tests/fpgadataflow/test_compilation_trafo.py b/tests/fpgadataflow/test_compilation_trafo.py index e1418b4febbfe1315198e19848dbf81a4e04658a..fd4174f27b2b9db1fae4f95589a5c9832679cc32 100644 --- a/tests/fpgadataflow/test_compilation_trafo.py +++ b/tests/fpgadataflow/test_compilation_trafo.py @@ -5,6 +5,7 @@ from onnx import TensorProto, helper import finn.core.utils as util from finn.core.datatype import DataType from finn.core.modelwrapper import ModelWrapper +from finn.transformation.fpgadataflow.cleanup import CleanUp from finn.transformation.fpgadataflow.codegen import CodeGen from finn.transformation.fpgadataflow.compile import Compile @@ -66,3 +67,4 @@ def test_compilation_trafo(): op type {} does not exist!""".format( node.op_type ) + model = model.transform(CleanUp()) diff --git a/tests/fpgadataflow/test_fpgadataflow_fclayer.py b/tests/fpgadataflow/test_fpgadataflow_fclayer.py index 22034834fbc365caf7f863cdd148916983d31153..4f1648feabd690da4a73b2d201dcaad4149725e7 100644 --- a/tests/fpgadataflow/test_fpgadataflow_fclayer.py +++ b/tests/fpgadataflow/test_fpgadataflow_fclayer.py @@ -8,6 +8,7 @@ import finn.custom_op.xnorpopcount as xp from finn.core.datatype import DataType from finn.core.modelwrapper import ModelWrapper from finn.core.utils import gen_finn_dt_tensor +from finn.transformation.fpgadataflow.cleanup import CleanUp from finn.transformation.fpgadataflow.codegen import CodeGen from finn.transformation.fpgadataflow.compile import Compile @@ -120,3 +121,4 @@ def test_fpgadataflow_fclayer_noact(idt, wdt, nf, sf, mw, mh): # execute model y_produced = oxe.execute_onnx(model, input_dict)["outp"] assert (y_produced.reshape(y_expected.shape) == y_expected).all() + model = model.transform(CleanUp()) diff --git a/tests/fpgadataflow/test_layer_streaming_maxpool_batch.py b/tests/fpgadataflow/test_layer_streaming_maxpool_batch.py index 2b69e9d8149d0f875c1b225f97aee01751a98f06..4a7ca1b5c0473c520f0e2ea775f7c8950eb16695 100644 --- a/tests/fpgadataflow/test_layer_streaming_maxpool_batch.py +++ b/tests/fpgadataflow/test_layer_streaming_maxpool_batch.py @@ -4,6 +4,7 @@ from onnx import TensorProto, helper import finn.core.onnx_exec as oxe from finn.core.datatype import DataType from finn.core.modelwrapper import ModelWrapper +from finn.transformation.fpgadataflow.cleanup import CleanUp from finn.transformation.fpgadataflow.codegen import CodeGen from finn.transformation.fpgadataflow.compile import Compile @@ -117,3 +118,4 @@ def test_layer_streaming_maxpool_batch(): input_dict = {"in": input_tensor} output_dict = oxe.execute_onnx(model, input_dict) print(output_dict) + model = model.transform(CleanUp())