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

[takipipo]

Refactor

• Remove code duplication from the review of Add support for loading and inferencing on aws-neuron and aws-neuronx ultralytics/ultralytics#18199 (review)

TODO

• Remove neuron_*.py
• Test exporting and loading NeuronX and Neuron models using YOLO

[takipipo]

@nirattisai-t @luangtatipsy Please consider this PR from the review of ultralytics#18199

[nirattisai-t]

LGTM

[takipipo]

@nirattisai-t could you test the following becuase this might effect the way kirin libraries use wisesight:ultralytics

TODO

• Remove neuron_*.py (ลบไฟล์ที่สร้างขึ้นมาออกให้หมด)
• Test exporting and loading NeuronX and Neuron models using YOLO (เรียก method ทุกอย่างผ่าน YOLO แทน)

something like this

from ultralytics import YOLO

model = YOLO("yolov8l.pt")
model.export(format="neuronx")
# model.export(format="neuron")

neuron_model = YOLO("yolov8l.neuronx")
# neuron_model = YOLO("yolov8l.neuron")

[nirattisai-t]

Okie, I will recheck those tests again if it works or not.

[nirattisai-t]

For neuronx there is an issue when trying to compile model with provided code
Environment

instance type :inf2.xlarge
image :ami-09e7c80512df5c25c

aws-neuronx-runtime-lib Version: 2.23.110.0-9b5179492
aws-neuronx-runtime-discovery==2.9
libneuronxla==2.1.681.0
neuronx-cc==2.16.345.0+69131dd3
neuronx-distributed==0.9.0
neuronx-distributed-training==1.0.1
torch-neuronx==2.5.1.2.4.0

Error Log

In [1]: from ultralytics import YOLO                                                                                
   ...: model = YOLO("yolov8l.pt")                                                                                  
   ...: model.export(format="neuronx")                                                                              
Ultralytics 8.3.49 🚀 Python-3.10.12 torch-2.5.1+cu124 CPU (unknown)                                                 
YOLOv8l summary (fused): 268 layers, 43,668,288 parameters, 0 gradients                                             
                                                                                                                    
PyTorch: starting from 'yolov8l.pt' with input shape (1, 3, 640, 640) BCHW and output shape(s) (1, 84, 8400) (83.7 M
B)                                                                                                                  
root: MASTER_ADDR environment variable is not set, defaulting to localhost                                          
root: Found libneuronpjrt.so. Setting PJRT_DEVICE=NEURON.                                                           
                                                                                                                    
AWS NeuronX: starting export with torch 2.5.1.2.4.0...                                                              
..                                                                                                                  
[GCA035]  Instruction: I-9647-0 with opcode: TensorTensor couldn't be allocated in SB                               
Memory Location Accessed:                                                                                           
multiply.3_deconcat_10613: 105624 Bytes per Partition and total of: 6759936 Bytes in SB                             
add.174.11160_i0: 640 Bytes per Partition and total of: 40960 Bytes in SB                                           
add.6: 109512 Bytes per Partition and total of: 7008768 Bytes in SB                                                 
Total Accessed Bytes per partition by instruction: 215776                                                           
Total SB Partition Size: 196608                                                                                     
 - Please open a support ticket at https://github.com/aws-neuron/aws-neuron-sdk/issues/new. You may also be able to 
obtain more information using the 'XLA_IR_DEBUG' and 'XLA_HLO_DEBUG' environment variables.                         
AWS NeuronX: export failure ❌ 39.1s: neuronx-cc failed with 70
---------------------------------------------------------------------------                                         
RuntimeError                              Traceback (most recent call last)                                         
Cell In[1], line 3                                                                                                  
      1 from ultralytics import YOLO                                                                                
      2 model = YOLO("yolov8l.pt")                                                                                  
----> 3 model.export(format="neuronx")                                                                              
                                                                                                                    
File ~/workspace/ultralytics/ultralytics/engine/model.py:738, in Model.export(self, **kwargs)                       
    730 custom = {                                                                                                  
    731     "imgsz": self.model.args["imgsz"],                                                                      
    732     "batch": 1,                                                                                             
   (...)
    735     "verbose": False,
    736 }  # method defaults
    737 args = {**self.overrides, **custom, **kwargs, "mode": "export"}  # highest priority args on the right
--> 738 return Exporter(overrides=args, _callbacks=self.callbacks)(model=self.model)

File ~/workspace/ultralytics/ultralytics/engine/exporter.py:392, in Exporter.__call__(self, model)
    390     f[13], _ = self.export_imx()
    391 if neuronx:  # NeuronX
--> 392     f[14], _ = self.export_neuronx()
    393 if neuron:  # Neuron
    394     f[15], _ = self.export_neuron()

File ~/workspace/ultralytics/ultralytics/engine/exporter.py:150, in try_export.<locals>.outer_func(*args, **kwargs)
    148 except Exception as e:
    149     LOGGER.error(f"{prefix} export failure ❌ {dt.t:.1f}s: {e}")
--> 150     raise e

File ~/workspace/ultralytics/ultralytics/engine/exporter.py:145, in try_export.<locals>.outer_func(*args, **kwargs)
    143 try:
    144     with Profile() as dt:
--> 145         f, model = inner_func(*args, **kwargs)
    146     LOGGER.info(f"{prefix} export success ✅ {dt.t:.1f}s, saved as '{f}' ({file_size(f):.1f} MB)")
    147     return f, model

File ~/workspace/ultralytics/ultralytics/engine/exporter.py:1253, in Exporter.export_neuronx(self, prefix)
   1251 LOGGER.info(f"\n{prefix} starting export with torch {torch_neuronx.__version__}...")
   1252 f = self.file.with_suffix(".neuronx")
-> 1253 ts = torch_neuronx.trace(self.model, self.im, strict=False)
   1254 extra_files = {"config.txt": json.dumps(self.metadata)}  # torch._C.ExtraFilesMap()
   1255 ts.save(str(f), _extra_files=extra_files)
File /opt/aws_neuronx_venv_pytorch/lib/python3.10/site-packages/torch_neuronx/xla_impl/trace.py:589, in trace(func, 
example_inputs, input_output_aliases, compiler_workdir, compiler_args, partitioner_config, inline_weights_to_neff, c
pu_backend, *_, **kwargs)                                                                                           
    584     return torch_neuronx.partition(                                                                         
    585         func, example_inputs, **(partitioner_config.__dict__)                                               
    586     )                                                                                                       
    588 with context, torch_neuronx.contexts.set_pjrt_device(cpu_backend):                                          
--> 589     neff_filename, metaneff, flattener, packer, weights = _trace(                                           
    590         func,                                                                                               
    591         example_inputs,                                                                                     
    592         states,
    593         input_output_aliases,
    594         compiler_workdir,
    595         compiler_args,                                                                                      
    596         inline_weights_to_neff,                                                                             
    597     )
    598     return create_neuron_model(                                                                             
    599         neff_filename,          
    600         metaneff,                                 
   (...)                                                  
    605         weights,    
    606     )                                             

File /opt/aws_neuronx_venv_pytorch/lib/python3.10/site-packages/torch_neuronx/xla_impl/trace.py:654, in _trace(func,
 example_inputs, states, input_output_aliases, compiler_workdir, compiler_args, inline_weights_to_neff)
    646 hlo_artifacts = generate_hlo(                                                                               
    647     func,  
    648     example_inputs,
    649     input_output_aliases=input_output_aliases,                                                              
    650     inline_weights_to_neff=inline_weights_to_neff,
    651 )                                                 
    653 # Call neuronx-cc to generate neff            
--> 654 neff_artifacts = generate_neff(                                                                             
    655     hlo_artifacts, 
    656     compiler_workdir=compiler_workdir,
    657     compiler_args=compiler_args,                                                                            
    658     inline_weights_to_neff=inline_weights_to_neff,                                                          
    659 )                                                 
    661 return (                                                                                                    
    662     neff_artifacts.neff_filename,                                                                           
    663     hlo_artifacts.metaneff.SerializeToString(),
   (...)
    666     hlo_artifacts.weights,         
    667 )                    
                                                                                                                    
File /opt/aws_neuronx_venv_pytorch/lib/python3.10/site-packages/torch_neuronx/xla_impl/trace.py:506, in generate_nef
f(hlo_artifacts, compiler_workdir, compiler_args, inline_weights_to_neff)
    498 compiler_target = setup_compiler_dirs(            
    499     hlo_artifacts.hlo_module,                                                                               
    500     compiler_workdir,                                                                                       
    501     hlo_artifacts.constant_parameter_tensors,                                                               
    502     inline_weights_to_neff,                                                                                 
    503 )                                                                                                           
    505 # Compile HLO to NEFF                             
--> 506 neff_filename = hlo_compile(                                                                                
    507     compiler_target,                                                                                        
    508     compiler_workdir,                             
    509     compiler_args,                                
    510 )                 
    512 return NeffArtifacts(neff_filename)   
                                                          
File /opt/aws_neuronx_venv_pytorch/lib/python3.10/site-packages/torch_neuronx/xla_impl/trace.py:396, in hlo_compile(
filename, compiler_workdir, compiler_args)
    389     elif status == -11:
    390         logger.warning(          
    391             "The neuronx-cc (neuron compiler) crashed (SEGFAULT). "
    392             "This is likely due to a bug in the compiler.  "
    393             "Please lodge an issue at 'https://github.com/aws/aws-neuron-sdk/issues'"
    394         )
--> 396     raise RuntimeError(f"neuronx-cc failed with {status}")
    398 return neff_filename                                                                                        
                                                                                                                    
RuntimeError: neuronx-cc failed with 70       

[nirattisai-t]

in neuron , from the provided code it was able to export the model but when i try to inference it cause an error because lib thought that i tried to load neuronx model.
Environment

instance type :inf1.xlarge
image :ami-09e7c80512df5c25c

aws-neuronx-runtime-discovery==2.9
libneuronxla==2.0.5347.0
neuron-cc==1.24.0.0+d58fa6134
neuronx-cc==2.15.143.0+e39249ad
neuronx-distributed==0.9.0
neuronx-distributed-training==1.0.1
torch==1.13.1
torch-neuron==1.13.1.2.11.13.0
torch-neuronx==2.1.2.2.3.2

In [3]: neuron_model.predict("../bus.jpg")                                                                          
---------------------------------------------------------------------------                                         
ImportError                               Traceback (most recent call last)                                         
Cell In[3], line 1                                                                                                  
----> 1 neuron_model.predict("../bus.jpg")                                                                          
                                                                                                                    
File ~/workspace/ultralytics/ultralytics/engine/model.py:551, in Model.predict(self, source, stream, predictor, **kw
args)                                                                                                               
    549 if not self.predictor:                                                                                      
    550     self.predictor = (predictor or self._smart_load("predictor"))(overrides=args, _callbacks=self.callbacks)
--> 551     self.predictor.setup_model(model=self.model, verbose=is_cli)                                            
    552 else:  # only update args if predictor is already setup                                                     
    553     self.predictor.args = get_cfg(self.predictor.args, args)                                                
                                                                                                                    
File ~/workspace/ultralytics/ultralytics/engine/predictor.py:308, in BasePredictor.setup_model(self, model, verbose)
    306 def setup_model(self, model, verbose=True):                                                                 
    307     """Initialize YOLO model with given parameters and set it to evaluation mode."""                        
--> 308     self.model = AutoBackend(                                                                               
    309         weights=model or self.args.model,                                                                   
    310         device=select_device(self.args.device, verbose=verbose),                                            
    311         dnn=self.args.dnn,                                                                                  
    312         data=self.args.data,                                                                                
    313         fp16=self.args.half,                                                                                
    314         batch=self.args.batch,                                                                              
    315         fuse=True,                                                                                          
    316         verbose=verbose,                                                                                    
    317     )                                                                                                       
    319     self.device = self.model.device  # update device                                                        
    320     self.args.half = self.model.fp16  # update half                                                         
                                                                                                                    
File /opt/aws_neuronx_venv_pytorch/lib/python3.10/site-packages/torch/autograd/grad_mode.py:27, in _DecoratorContext
Manager.__call__.<locals>.decorate_context(*args, **kwargs)                                                         
     24 @functools.wraps(func)
     25 def decorate_context(*args, **kwargs):
     26     with self.clone():
---> 27         return func(*args, **kwargs)

File ~/workspace/ultralytics/ultralytics/nn/autobackend.py:185, in AutoBackend.__init__(self, weights, device, dnn, 
data, fp16, batch, fuse, verbose)
    183 # NeuronX
    184 elif neuronx:
--> 185     import torch_neuronx
    187     LOGGER.info(f"Loading {w} for NeuronX version {torch_neuronx.__version__} inference... ")
    188     extra_files = {"config.txt": ""}  # model metadata
�
File /opt/aws_neuronx_venv_pytorch/lib/python3.10/site-packages/torch_neuronx/__init__.py:9
      7 from . import contexts
      8 from . import experimental
----> 9 from .xla_impl.bucket_trace import bucket_model_trace, BucketModel, BucketModelScript, BucketModelConfig
     10 from .xla_impl.trace import trace, move_trace_to_device
     11 from .xla_impl.options import Options

File /opt/aws_neuronx_venv_pytorch/lib/python3.10/site-packages/torch_neuronx/xla_impl/bucket_trace.py:10
      7 import torch_neuronx
      8 from torch_neuronx.xla_impl import structure
---> 10 import torch_xla.distributed.xla_multiprocessing as xmp
     11 from torch_xla.utils.utils import get_free_tcp_ports
     14 def identity(
     15     shapes_collection: List[List[List[int]]],
     16     states: List[torch.Tensor],
     17     bucket_idx_tensor: torch.Tensor,
     18 ):

File /opt/aws_neuronx_venv_pytorch/lib/python3.10/site-packages/torch_xla/__init__.py:112
    109 from ._patched_functions import _apply_patches
    110 from .version import __version__
--> 112 import _XLAC
    114 _found_libtpu = _setup_tpu_vm_library_path()
    116 # Setup Neuron library for AWS EC2 inf/trn instances.

ImportError: /opt/aws_neuronx_venv_pytorch/lib/python3.10/site-packages/_XLAC.cpython-310-x86_64-linux-gnu.so: undef
ined symbol: _ZNK5torch4lazy17LazyGraphExecutor16ShouldSyncTensorERKN3c1013intrusive_ptrINS0_10LazyTensorENS2_6detai
l34intrusive_target_default_null_typeIS4_EEEE

Even i replaced neuronx section with neuron section code to run prediction, their still have an error as shown below

In [3]: neuron_model.predict("../bus.jpg")                                                                          
Loading yolov8l.neuron for Neuron version 1.13.1.2.11.13.0 inference... 
                             
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)                                         
Cell In[3], line 1
----> 1 neuron_model.predict("../bus.jpg")
                                                                                                                    
File ~/workspace/ultralytics/ultralytics/engine/model.py:558, in Model.predict(self, source, stream, predictor, **kw
args)                                                                                                               
    556 if prompts and hasattr(self.predictor, "set_prompts"):  # for SAM-type models
    557     self.predictor.set_prompts(prompts)                                                                     
--> 558 return self.predictor.predict_cli(source=source) if is_cli else self.predictor(source=source, stream=stream)
                                                          
File ~/workspace/ultralytics/ultralytics/engine/predictor.py:173, in BasePredictor.__call__(self, source, model, str
eam, *args, **kwargs)                                                                                               
    171     return self.stream_inference(source, model, *args, **kwargs)
    172 else:                                                                                                       
--> 173     return list(self.stream_inference(source, model, *args, **kwargs))
                                                          
File /opt/aws_neuronx_venv_pytorch/lib/python3.10/site-packages/torch/autograd/grad_mode.py:43, in _DecoratorContext
Manager._wrap_generator.<locals>.generator_context(*args, **kwargs)
     40 try:                                                                                                        
     41     # Issuing `None` to a generator fires it up
     42     with self.clone():
---> 43         response = gen.send(None)                                                                           
     45     while True:
     46         try:                                      
     47             # Forward the response to our caller and get its next request
                                                                                                                    
File ~/workspace/ultralytics/ultralytics/engine/predictor.py:259, in BasePredictor.stream_inference(self, source, mo
del, *args, **kwargs)
    257 # Inference                                                                                                 
    258 with profilers[1]:
--> 259     preds = self.inference(im, *args, **kwargs)
    260     if self.args.embed:                           
    261         yield from [preds] if isinstance(preds, torch.Tensor) else preds  # yield embedding tensors         

File ~/workspace/ultralytics/ultralytics/engine/predictor.py:143, in BasePredictor.inference(self, im, *args, **kwar
gs)                                                                                                                 
    137 """Runs inference on a given image using the specified model and arguments."""                              
    138 visualize = (        
    139     increment_path(self.save_dir / Path(self.batch[0][0]).stem, mkdir=True)                                 
    140     if self.args.visualize and (not self.source_type.tensor)    
    141     else False       
    142 )                                                                                                           
--> 143 return self.model(im, augment=self.args.augment, visualize=visualize, embed=self.args.embed, *args, **kwargs
)                 
                                                          
File /opt/aws_neuronx_venv_pytorch/lib/python3.10/site-packages/torch/nn/modules/module.py:1194, in Module._call_imp
l(self, *input, **kwargs)                                                                                           
   1190 # If we don't have any hooks, we want to skip the rest of the logic in                                      
   1191 # this function, and just call forward.                                                                     
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks    
   1193         or _global_forward_hooks or _global_forward_pre_hooks):                                             
-> 1194     return forward_call(*input, **kwargs)         
   1195 # Do not call functions when jit is used                                                                    
   1196 full_backward_hooks, non_full_backward_hooks = [], []                                                       
                                                                                                                    
File ~/workspace/ultralytics/ultralytics/nn/autobackend.py:686, in AutoBackend.forward(self, im, augment, visualize,
 embed)                                                                                                             
    684     y = self.frozen_func(x=self.tf.constant(im))  
    685 else:  # Lite or Edge TPU                                                                                   
--> 686     details = self.input_details[0]                                                                         
    687     is_int = details["dtype"] in {np.int8, np.int16}  # is TFLite quantized int8 or int16 model             
    688     if is_int:                                    
                                                          
File /opt/aws_neuronx_venv_pytorch/lib/python3.10/site-packages/torch/nn/modules/module.py:1269, in Module.__getattr
__(self, name)         
   1267     if name in modules:                           
   1268         return modules[name]                                                                                
-> 1269 raise AttributeError("'{}' object has no attribute '{}'".format(                                            
   1270     type(self).__name__, name))                                                                             
                             
AttributeError: 'AutoBackend' object has no attribute 'input_details'