refactor: remove code duplication but still keeping neuron_*.py

ultralytics/engine/exporter.py

@@ -19,6 +19,8 @@
 MNN                     | `mnn`                     | yolo11n.mnn
 NCNN                    | `ncnn`                    | yolo11n_ncnn_model/
 IMX                     | `imx`                     | yolo11n_imx_model/
+AWS NeuronX             | `neuronx`                 | yolo11n.neuronx
+AWS Neuron              | `neuron`                  | yolo11n.neuron
 
 Requirements:
     $ pip install "ultralytics[export]"
@@ -116,6 +118,8 @@ def export_formats():
         ["MNN", "mnn", ".mnn", True, True],
         ["NCNN", "ncnn", "_ncnn_model", True, True],
         ["IMX", "imx", "_imx_model", True, True],
+        ["AWS NeuronX", "neuronx", ".neuronx", True, True],
+        ["AWS Neuron", "neuron", ".neuron", True, True],
     ]
     return dict(zip(["Format", "Argument", "Suffix", "CPU", "GPU"], zip(*x)))
 
@@ -210,6 +214,8 @@ def __call__(self, model=None) -> str:
             mnn,
             ncnn,
             imx,
+            neuronx,
+            neuron,
         ) = flags  # export booleans
         is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs))
 
@@ -382,6 +388,10 @@ def __call__(self, model=None) -> str:
             f[12], _ = self.export_ncnn()
         if imx:
             f[13], _ = self.export_imx()
+        if neuronx:  # NeuronX
+            f[14], _ = self.export_neuronx()
+        if neuron:  # Neuron
+            f[15], _ = self.export_neuron()
 
         # Finish
         f = [str(x) for x in f if x]  # filter out '' and None
@@ -1232,6 +1242,30 @@ def forward(self, images):
             file.writelines([f"{name}\n" for _, name in self.model.names.items()])
 
         return f, None
+
+    @try_export
+    def export_neuronx(self, prefix=colorstr("AWS NeuronX:")):
+        import torch_neuronx
+
+        """YOLOv8 NeuronX model export."""
+        LOGGER.info(f"\n{prefix} starting export with torch {torch_neuronx.__version__}...")
+        f = self.file.with_suffix(".neuronx")
+        ts = torch_neuronx.trace(self.model, self.im, strict=False)
+        extra_files = {"config.txt": json.dumps(self.metadata)}  # torch._C.ExtraFilesMap()
+        ts.save(str(f), _extra_files=extra_files)
+        return f, None
+
+    @try_export
+    def export_neuron(self, prefix=colorstr("AWS Neuron:")):
+        import torch_neuron
+
+        """YOLOv8 Neuron model export."""
+        LOGGER.info(f"\n{prefix} starting export with torch {torch_neuron.__version__}...")
+        f = self.file.with_suffix(".neuron")
+        ts = torch_neuron.trace(self.model, self.im, strict=False)
+        extra_files = {"config.txt": json.dumps(self.metadata)}
+        ts.save(str(f), _extra_files=extra_files)
+        return f, None
 
     def _add_tflite_metadata(self, file):
         """Add metadata to *.tflite models per https://www.tensorflow.org/lite/models/convert/metadata."""

ultralytics/nn/autobackend.py

@@ -75,6 +75,8 @@ class AutoBackend(nn.Module):
             | PaddlePaddle          | *_paddle_model/   |
             | MNN                   | *.mnn             |
             | NCNN                  | *_ncnn_model/     |
+            | AWS NeuronX           | *.neuronx         |
+            | AWS Neuron            | *.neuron          |
 
     This class offers dynamic backend switching capabilities based on the input model format, making it easier to deploy
     models across various platforms.
@@ -125,6 +127,8 @@ def __init__(
             ncnn,
             imx,
             triton,
+            neuronx,
+            neuron,
         ) = self._model_type(w)
         fp16 &= pt or jit or onnx or xml or engine or nn_module or triton  # FP16
         nhwc = coreml or saved_model or pb or tflite or edgetpu  # BHWC formats (vs torch BCWH)
@@ -176,7 +180,26 @@ def __init__(
             model.half() if fp16 else model.float()
             if extra_files["config.txt"]:  # load metadata dict
                 metadata = json.loads(extra_files["config.txt"], object_hook=lambda x: dict(x.items()))
+        # NeuronX
+        elif neuronx:
+            import torch_neuronx
 
+            LOGGER.info(f"Loading {w} for NeuronX version {torch_neuronx.__version__} inference... ")
+            extra_files = {"config.txt": ""}  # model metadata
+            model = torch.jit.load(w, _extra_files=extra_files, map_location=device)
+            model.half() if fp16 else model.float()
+            if extra_files["config.txt"]:  # load metadata dict
+                metadata = json.loads(extra_files["config.txt"], object_hook=lambda x: dict(x.items()))
+        ## Neuron
+        elif neuron:
+            import torch_neuron
+
+            LOGGER.info(f"Loading {w} for Neuron version {torch_neuron.__version__} inference... ")
+            extra_files = {"config.txt": ""}  # model metadata
+            model = torch.jit.load(w, _extra_files=extra_files, map_location=device)
+            model.half() if fp16 else model.float()
+            if extra_files["config.txt"]:  # load metadata dict
+                metadata = json.loads(extra_files["config.txt"], object_hook=lambda x: dict(x.items()))
         # ONNX OpenCV DNN
         elif dnn:
             LOGGER.info(f"Loading {w} for ONNX OpenCV DNN inference...")