Feature/support inf1

ultralytics/engine/neuron_exporter.py

@@ -105,6 +105,7 @@ def export_formats():
         ["PyTorch", "-", ".pt", True, True],
         ["TorchScript", "torchscript", ".torchscript", True, True],
         ["AWS NeuronX", "neuronx", ".neuronx", True, True],
+        ["AWS Neuron", "neuron", ".neuron", True, True],
         ["ONNX", "onnx", ".onnx", True, True],
         ["OpenVINO", "openvino", "_openvino_model", True, False],
         ["TensorRT", "engine", ".engine", False, True],
@@ -173,7 +174,22 @@ def __call__(self, model=None) -> str:
         flags = [x == fmt for x in fmts]
         if sum(flags) != 1:
             raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
-        jit, neuronx, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, ncnn = flags  # export booleans
+        (
+            jit,
+            neuronx,
+            neuron,
+            onnx,
+            xml,
+            engine,
+            coreml,
+            saved_model,
+            pb,
+            tflite,
+            edgetpu,
+            tfjs,
+            paddle,
+            ncnn,
+        ) = flags  # export booleans
         is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs))
 
         # Device
@@ -315,6 +331,8 @@ def __call__(self, model=None) -> str:
             f[11], _ = self.export_ncnn()
         if neuronx:  # NeuronX
             f[12], _ = self.export_neuronx()
+        if neuron:  # Neuron
+            f[13], _ = self.export_neuron()
 
         # Finish
         f = [str(x) for x in f if x]  # filter out '' and None
@@ -354,6 +372,17 @@ def export_neuronx(self, prefix=colorstr("AWS NeuronX:")):
             from torch.utils.mobile_optimizer import optimize_for_mobile
 
             optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
-        else:
-            ts.save(str(f), _extra_files=extra_files)
-        return f, None
+        extra_files = {"config.txt": json.dumps(self.metadata)}
+        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

ultralytics/nn/neuron_autobackend.py

@@ -110,6 +110,7 @@ def __init__(
             pt,
             jit,
             neuronx,
+            neuron,
             onnx,
             xml,
             engine,
@@ -182,6 +183,16 @@ 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()))
+        ## Neuron
+        elif neuron:
+            import torch_neuron
+            LOGGER.info(f"Loading {w} for Neuron 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...")
@@ -469,6 +480,8 @@ def forward(self, im, augment=False, visualize=False, embed=None):
 
         elif self.neuronx:
             y = self.model(im)
+        elif self.neuron:
+            y = self.model(im)
         # ONNX OpenCV DNN
         elif self.dnn:
             im = im.cpu().numpy()  # torch to numpy