diff --git a/aidge_export_arm_cortexm/operators.py b/aidge_export_arm_cortexm/operators.py index 24611dde638109bcd19790870cf73c561e46601f..b4e6aa85cd8709dd0944c4eaeb2b25560ab8b945 100644 --- a/aidge_export_arm_cortexm/operators.py +++ b/aidge_export_arm_cortexm/operators.py @@ -51,9 +51,25 @@ class Producer_ARMCortexM(ExportNode): def __init__(self, node, mem_info, is_input, is_output): super().__init__(node, mem_info, is_input, is_output) self.values = np.array(self.operator.get_output(0)) - + print(f"{node.name()}: {self.values.shape}") if len(self.values.shape) == 4: # Note: export in HWC self.values = np.transpose(self.values, (0, 2, 3, 1)) + # The following block of code is a dirty fix for FC + # The issue is that FC weight in Aidge are made for an CHW input + # Current export is made with HWC format + # So we need to reorder weights of the FC + # Note: it is not necessary if H and W != 1 (equivalent to in_dims length == 4) + + if len(self.values.shape) == 2: + parents = node.get_children() + if len(parents) == 1 and list(parents)[0].type() == "FC": + data_in = list(parents)[0].get_operator().get_input(0) + if len(data_in.dims()) == 4: + C = data_in.dims()[1] + H = data_in.dims()[2] + W = data_in.dims()[3] + # Transpose weights to adapt the HWC + self.values = self.values.reshape(-1, C, H, W).transpose(0, 2, 3, 1) def export(self, export_folder: Path): header_path = f"include/parameters/{self.attributes['name']}.hpp"