diff --git a/aidge_export_cpp/operators/Conv.py b/aidge_export_cpp/operators/Conv.py
index a813c2f6c82743b0e8983894dab62450eeb8d7c9..88fbf3bda71ec9b7589f4faec21a9386126f7e55 100644
--- a/aidge_export_cpp/operators/Conv.py
+++ b/aidge_export_cpp/operators/Conv.py
@@ -1,6 +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
@ExportLibCpp.register("Conv1D",
aidge_core.ImplSpec(
@@ -136,7 +137,8 @@ class QConv(Conv):
[ # Output specifications
aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
],
- ))
+ ),
+ PaddedInPlace_CP.default_model)
class PadConv(QConv):
def __init__(self, node, mem_info):
super().__init__(node, mem_info)
@@ -178,7 +180,8 @@ class ConvAct(QConv):
[ # Output specifications
aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
],
- ))
+ ),
+ PaddedInPlace_CP.default_model)
class PadConvAct(PadConv, ConvAct):
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 a38030e362a015f520119e6322fb230d8599764d..d0587f545182005a4fe8c827de468fff3f94cb11 100644
--- a/aidge_export_cpp/operators/ConvDw.py
+++ b/aidge_export_cpp/operators/ConvDw.py
@@ -1,6 +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
@ExportLibCpp.register("ConvDepthWise2D",
aidge_core.ImplSpec(
@@ -87,7 +88,8 @@ class QConvDw(ConvDw):
[ # Output specifications
aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
],
- ))
+ ),
+ PaddedInPlace_CP.default_model)
class PadConvDw(QConvDw):
def __init__(self, node, mem_info):
super().__init__(node, mem_info)
@@ -129,7 +131,8 @@ class ConvDwAct(QConvDw):
[ # Output specifications
aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
],
- ))
+ ),
+ PaddedInPlace_CP.default_model)
class PadConvDwAct(PadConvDw, ConvDwAct):
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 f84f2cff0f1df283327c1122ddb7b77049bddd3f..57fd86d27243678f2162a3ac845d21424a801f8d 100644
--- a/aidge_export_cpp/operators/Pad.py
+++ b/aidge_export_cpp/operators/Pad.py
@@ -3,6 +3,28 @@ 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):
+ 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)
+
+ pad_node = self.get_operator().get_micro_graph().get_ordered_inputs()[0][0]
+ input = pad_node.get_operator().get_input(0)
+ if input:
+ output = pad_node.get_operator().get_output(0)
+ return aidge_core.Elts_t.data_elts(output.size() - input.size())
+ else:
+ return aidge_core.Elts_t.none_elts()
+
+
@ExportLibCpp.register("Pad2D", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any)))
class CppPad(ExportNodeCpp):
def __init__(self, node, mem_info):
diff --git a/aidge_export_cpp/operators/Pool.py b/aidge_export_cpp/operators/Pool.py
index d6e9bfc82b4db19d9a800928c728e4c0209aa61a..9c8df131d79faf5fccce4d07e8d022235f34cacc 100644
--- a/aidge_export_cpp/operators/Pool.py
+++ b/aidge_export_cpp/operators/Pool.py
@@ -2,6 +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
class Pool(ExportNodeCpp):
def __init__(self, node, mem_info):
@@ -66,7 +67,9 @@ class MaxPool(Pool):
self.attributes["stride_dims"] = PoolNode[0].get_operator().attr.stride_dims
-@ExportLibCpp.register_metaop(["PaddedMaxPooling2D", "PadMaxPool"], aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)))
+@ExportLibCpp.register_metaop(["PaddedMaxPooling2D", "PadMaxPool"],
+ aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
+ PaddedInPlace_CP.default_model)
class PadMaxPool(MaxPool, PadPool):
def __init__(self, node, mem_info):
super().__init__(node, mem_info)
@@ -80,7 +83,9 @@ class MaxPoolAct(MaxPool, PoolAct):
super().__init__(node, mem_info)
-@ExportLibCpp.register_metaop("PadMaxPoolAct", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)))
+@ExportLibCpp.register_metaop("PadMaxPoolAct",
+ aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
+ PaddedInPlace_CP.default_model)
class PadMaxPoolAct(PadMaxPool, MaxPoolAct):
def __init__(self, node, mem_info):
super().__init__(node, mem_info)
@@ -100,7 +105,9 @@ class AvgPool(Pool):
self.attributes["stride_dims"] = PoolNode[0].get_operator().attr.stride_dims
-@ExportLibCpp.register_metaop(["PaddedAvgPooling2D", "PadAvgPool"], aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)))
+@ExportLibCpp.register_metaop(["PaddedAvgPooling2D", "PadAvgPool"],
+ aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
+ PaddedInPlace_CP.default_model)
class PadAvgPool(AvgPool, PadPool):
def __init__(self, node, mem_info):
super().__init__(node, mem_info)
@@ -114,7 +121,9 @@ class AvgPoolAct(AvgPool, PoolAct):
super().__init__(node, mem_info)
-@ExportLibCpp.register_metaop("PadAvgPoolAct", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)))
+@ExportLibCpp.register_metaop("PadAvgPoolAct",
+ aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
+ PaddedInPlace_CP.default_model)
class PadAvgPoolAct(PadAvgPool, AvgPoolAct):
def __init__(self, node, mem_info):
super().__init__(node, mem_info)
@@ -131,7 +140,9 @@ class GlobalAvgPool(Pool):
self.attributes["kernel_dims"] = [self.attributes["in_width"][0], self.attributes["in_height"][0]]
-@ExportLibCpp.register_metaop("PadGlobalAvgPool", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)))
+@ExportLibCpp.register_metaop("PadGlobalAvgPool",
+ aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
+ PaddedInPlace_CP.default_model)
class PadGlobalAvgPool(GlobalAvgPool, PadPool):
def __init__(self, node, mem_info):
super().__init__(node, mem_info)
@@ -145,7 +156,9 @@ class GlobalAvgPoolAct(GlobalAvgPool, PoolAct):
super().__init__(node, mem_info)
-@ExportLibCpp.register_metaop("PadGlobalAvgPoolAct", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)))
+@ExportLibCpp.register_metaop("PadGlobalAvgPoolAct",
+ aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
+ PaddedInPlace_CP.default_model)
class PadGlobalAvgPoolAct(PadGlobalAvgPool, GlobalAvgPoolAct):
def __init__(self, node, mem_info):
super().__init__(node, mem_info)
\ No newline at end of file
diff --git a/examples/export_ResNet18/resnet18.py b/examples/export_ResNet18/resnet18.py
index b27cafd7e75da9e16392e4c5844c93f36bd3a86a..4cca5c8c58cba71b6e19d972d28ce76d14933e79 100644
--- a/examples/export_ResNet18/resnet18.py
+++ b/examples/export_ResNet18/resnet18.py
@@ -42,6 +42,7 @@ supported_types = ["float32", "int8"]
parser = argparse.ArgumentParser(description="Export the ResNet18 model with the aidge_export_cpp module.")
parser.add_argument("--dev", action="store_true", help="Export in dev mode")
+parser.add_argument("--mem_wrap", action="store_true", help="Use memory wrapping")
parser.add_argument("--no_cuda", action="store_true", help="Disable USE_CUDA usage to perform inferences and training.")
parser.add_argument("--dtype", type=str, choices=supported_types, default="float32", help="Specify the targeted datatype : [int8, float32]")
parser.add_argument("--aidge_cmp", action="store_true", help="Use aidge tensor results as reference.")
@@ -559,7 +560,8 @@ aidge_export_cpp.export(EXPORT_FOLDER,
labels = aidge_core.Tensor(labels[0]),
#inputs_tensor=inputs_tensor,
dev_mode = DEV_MODE,
- aidge_cmp = AIDGE_CMP)
+ aidge_cmp = AIDGE_CMP,
+ memory_manager_args = {"wrapping": True} if args.mem_wrap else {})
print("\n### Compiling the export ###")
try: