diff --git a/aidge_export_cpp/operators/Conv.py b/aidge_export_cpp/operators/Conv.py index 676c0168181193f22ae6a22c14dfa99e3f39486c..44b50b49a2ecf4c5957fc842aee01a19c3e0ecb1 100644 --- a/aidge_export_cpp/operators/Conv.py +++ b/aidge_export_cpp/operators/Conv.py @@ -1,7 +1,53 @@ import aidge_core from aidge_core.export_utils import ExportNodeCpp, get_node_from_metaop from aidge_export_cpp import ROOT, ExportLibCpp, set_scaling_attributes -from .Pad import PaddedInPlace_CP + +# Consumer-Producer model to allow memory wrapping for Conv/PaddedConv +# (and Pool/PaddedPool), keeping one input line margin in NHWC data format +# (one input line = W*C) +class PaddedInPlace_CP(aidge_core.ProdConso): + def __init__(self, op: aidge_core.Operator): + aidge_core.ProdConso.__init__(self, op, False) + + def default_model(op: aidge_core.Operator): + return PaddedInPlace_CP(op) + + def get_nb_required_protected(self, input_idx): + if input_idx != 0: + return super().get_nb_required_protected(input_idx) + + input = self.get_operator().get_input(0) + if not input: + return aidge_core.Elts_t.none_elts() + + # Non-Padded case: margin = one input line + margin = 1 + if not self.get_operator().is_atomic(): + # Padded case: margin = (padding_y / stride_y) input lines + sub_graph = self.get_operator().get_micro_graph().clone() + aidge_core.expand_metaops(sub_graph, True) + + padding_y = 0 + stride_y = 1 + for node in sub_graph.get_nodes(): + if hasattr(node.get_operator().attr, 'stride_dims'): + if len(node.get_operator().attr.stride_dims) > 1: + stride_y = node.get_operator().attr.stride_dims[0] + elif hasattr(node.get_operator().attr, 'begin_end_borders'): + if len(node.get_operator().attr.begin_end_borders) > 2: + padding_y = node.get_operator().attr.begin_end_borders[0] + + margin += padding_y // stride_y + + if len(input.dims()) == 4: + # 2D: one input line = W*C + margin *= input.dims()[2] * input.dims()[3] + else: + # 1D: one input line = C + margin *= input.dims()[2] + + return aidge_core.Elts_t.data_elts(margin) + @ExportLibCpp.register("Conv1D", aidge_core.ImplSpec( @@ -14,7 +60,7 @@ from .Pad import PaddedInPlace_CP aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nwc) ], ), - aidge_core.ProdConso.in_place_model) + PaddedInPlace_CP.default_model) class Conv1D(ExportNodeCpp): def __init__(self, node, mem_info): super().__init__(node, mem_info) @@ -64,7 +110,7 @@ class Conv1D(ExportNodeCpp): aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc) ], ), - aidge_core.ProdConso.in_place_model) + PaddedInPlace_CP.default_model) class Conv2D(ExportNodeCpp): def __init__(self, node, mem_info): super().__init__(node, mem_info) @@ -114,7 +160,7 @@ class Conv2D(ExportNodeCpp): aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc) ], ), - aidge_core.ProdConso.in_place_model) + PaddedInPlace_CP.default_model) class QConv(Conv2D): def __init__(self, node, mem_info): super().__init__(node, mem_info) @@ -159,7 +205,7 @@ class PadConv(QConv): aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc) ], ), - aidge_core.ProdConso.in_place_model) + PaddedInPlace_CP.default_model) class ConvAct(QConv): def __init__(self, node, mem_info): super().__init__(node, mem_info) diff --git a/aidge_export_cpp/operators/ConvDw.py b/aidge_export_cpp/operators/ConvDw.py index d0587f545182005a4fe8c827de468fff3f94cb11..136666719552aa58ac8e7a08c363c4cfd128c5e2 100644 --- a/aidge_export_cpp/operators/ConvDw.py +++ b/aidge_export_cpp/operators/ConvDw.py @@ -1,7 +1,7 @@ import aidge_core from aidge_core.export_utils import ExportNodeCpp, get_node_from_metaop from aidge_export_cpp import ROOT, ExportLibCpp, set_scaling_attributes -from .Pad import PaddedInPlace_CP +from .Conv import PaddedInPlace_CP @ExportLibCpp.register("ConvDepthWise2D", aidge_core.ImplSpec( @@ -14,7 +14,7 @@ from .Pad import PaddedInPlace_CP aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc) ], ), - aidge_core.ProdConso.in_place_model) + PaddedInPlace_CP.default_model) class ConvDw(ExportNodeCpp): def __init__(self, node, mem_info): super().__init__(node, mem_info) @@ -65,7 +65,7 @@ class ConvDw(ExportNodeCpp): aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc) ], ), - aidge_core.ProdConso.in_place_model) + PaddedInPlace_CP.default_model) class QConvDw(ConvDw): def __init__(self, node, mem_info): super().__init__(node, mem_info) @@ -110,7 +110,7 @@ class PadConvDw(QConvDw): aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc) ], ), - aidge_core.ProdConso.in_place_model) + PaddedInPlace_CP.default_model) class ConvDwAct(QConvDw): def __init__(self, node, mem_info): super().__init__(node, mem_info) diff --git a/aidge_export_cpp/operators/Pad.py b/aidge_export_cpp/operators/Pad.py index 05aeb2b38b6bc37bc4b7ecf3dc0ac121db60b81e..bc73ef43ab862bf11a2a41477006827ffcb19afa 100644 --- a/aidge_export_cpp/operators/Pad.py +++ b/aidge_export_cpp/operators/Pad.py @@ -3,20 +3,19 @@ from aidge_core.export_utils import ExportNodeCpp from aidge_export_cpp import ROOT from aidge_export_cpp import ExportLibCpp -# Consumer-Producer model to allow memory wrapping for Padded in-place operators -# like PaddedConv ou PaddedMax/AvgPool -class PaddedInPlace_CP(aidge_core.ProdConso): +# Consumer-Producer model to allow memory wrapping for Pad in-place operator +class PadInPlace_CP(aidge_core.ProdConso): def __init__(self, op: aidge_core.Operator): aidge_core.ProdConso.__init__(self, op, False) def default_model(op: aidge_core.Operator): - return PaddedInPlace_CP(op) + return PadInPlace_CP(op) def get_nb_required_protected(self, input_idx): if input_idx != 0: return super().get_nb_required_protected(input_idx) - pad_node = self.get_operator().get_micro_graph().get_ordered_inputs()[0][0] + pad_node = self.get_operator() input = pad_node.get_operator().get_input(0) if input: output = pad_node.get_operator().get_output(0) @@ -25,7 +24,9 @@ class PaddedInPlace_CP(aidge_core.ProdConso): return aidge_core.Elts_t.none_elts() -@ExportLibCpp.register("Pad2D", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any))) +@ExportLibCpp.register("Pad2D", + aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any)), + PadInPlace_CP.default_model) class CppPad(ExportNodeCpp): def __init__(self, node, mem_info): super().__init__(node, mem_info) diff --git a/aidge_export_cpp/operators/Pool.py b/aidge_export_cpp/operators/Pool.py index 9c8df131d79faf5fccce4d07e8d022235f34cacc..1d6cd53d6d0f2d7c000b8f3acc81ee5771295080 100644 --- a/aidge_export_cpp/operators/Pool.py +++ b/aidge_export_cpp/operators/Pool.py @@ -2,7 +2,7 @@ import aidge_core from aidge_core.export_utils import ExportNodeCpp, get_node_from_metaop from aidge_export_cpp import ROOT from aidge_export_cpp import ExportLibCpp -from .Pad import PaddedInPlace_CP +from .Conv import PaddedInPlace_CP class Pool(ExportNodeCpp): def __init__(self, node, mem_info): @@ -55,7 +55,7 @@ class PoolAct(Pool): @ExportLibCpp.register("MaxPooling2D", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)), - aidge_core.ProdConso.in_place_model) + PaddedInPlace_CP.default_model) class MaxPool(Pool): def __init__(self, node, mem_info): super().__init__(node, mem_info) @@ -77,7 +77,7 @@ class PadMaxPool(MaxPool, PadPool): @ExportLibCpp.register_metaop("MaxPoolAct", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)), - aidge_core.ProdConso.in_place_model) + PaddedInPlace_CP.default_model) class MaxPoolAct(MaxPool, PoolAct): def __init__(self, node, mem_info): super().__init__(node, mem_info) @@ -93,7 +93,7 @@ class PadMaxPoolAct(PadMaxPool, MaxPoolAct): @ExportLibCpp.register("AvgPooling2D", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)), - aidge_core.ProdConso.in_place_model) + PaddedInPlace_CP.default_model) class AvgPool(Pool): def __init__(self, node, mem_info): super().__init__(node, mem_info) @@ -115,7 +115,7 @@ class PadAvgPool(AvgPool, PadPool): @ExportLibCpp.register_metaop("AvgPoolAct", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)), - aidge_core.ProdConso.in_place_model) + PaddedInPlace_CP.default_model) class AvgPoolAct(AvgPool, PoolAct): def __init__(self, node, mem_info): super().__init__(node, mem_info)