Cannot export wraped YOLO model to onnx

I want to warp the YOLOv11 model as a torch.nn.Module and as ROI detection module, another torch.nn.Module would take over the output from the ROI detection module. The prediction runs well.

# code to warp YOLO
class ROIDetectionModule(nn.Module):
    """This model will perform object detection on input images and crop the images based on
    the detection results.
    """
    def __init__(self, 
                 model_cpt_path: str = None,
                 DEBUG: bool = False,
                 threshold: float = 0.6):
        super(ROIDetectionModule, self).__init__()        
        try:
            self.model = YOLO(model_cpt_path)
        except Exception as e:
            raise Exception(f"Failed to load model from checkpoint file {model_cpt_path}. Error: {e}")
        self.DEBUG = DEBUG
        self.threshold = threshold

However, when I trying to export my combined Module to onnx, it trigerd infinity training of YOLO. The generated training folders are ‘train12’, ‘train122’, ‘train1222’, …

# code the export module
torch.onnx.export(my_combined_module, 
                        dummy_input, 
                        "multimodel.onnx", 
                        verbose=True, 
                        dynamic_axes={'input': { 2: 'height', 3: 'width'}})

By the way, I found model.eval() would triger infinity training as well.

The Error returned is here.

  File "d:\usr\Projects\floorplan-dev\experiments\floorplan\deep_learning_exp\nn2onnx.py", line 64, in main
    torch.onnx.export(mhmtm,
  File "C:\Users\MH_3002\anaconda3\envs\multitask-floorplan-3.10\lib\site-packages\torch\onnx\utils.py", line 551, in export
    _export(
  File "C:\Users\MH_3002\anaconda3\envs\multitask-floorplan-3.10\lib\site-packages\torch\onnx\utils.py", line 1625, in _export
    with exporter_context(model, training, verbose):
  File "C:\Users\MH_3002\anaconda3\envs\multitask-floorplan-3.10\lib\contextlib.py", line 135, in __enter__
    return next(self.gen)
  File "C:\Users\MH_3002\anaconda3\envs\multitask-floorplan-3.10\lib\site-packages\torch\onnx\utils.py", line 180, in exporter_context
    with select_model_mode_for_export(
  File "C:\Users\MH_3002\anaconda3\envs\multitask-floorplan-3.10\lib\contextlib.py", line 135, in __enter__
    return next(self.gen)
  File "C:\Users\MH_3002\anaconda3\envs\multitask-floorplan-3.10\lib\site-packages\torch\onnx\utils.py", line 120, in select_model_mode_for_export
    model.train(False)
  File "C:\Users\MH_3002\anaconda3\envs\multitask-floorplan-3.10\lib\site-packages\torch\nn\modules\module.py", line 2456, in train
    module.train(mode)
  File "C:\Users\MH_3002\anaconda3\envs\multitask-floorplan-3.10\lib\site-packages\torch\nn\modules\module.py", line 2456, in train
    module.train(mode)
  File "C:\Users\MH_3002\anaconda3\envs\multitask-floorplan-3.10\lib\site-packages\ultralytics\engine\model.py", line 796, in train
    self.trainer = (trainer or self._smart_load("trainer"))(overrides=args, _callbacks=self.callbacks)
  File "C:\Users\MH_3002\anaconda3\envs\multitask-floorplan-3.10\lib\site-packages\ultralytics\engine\trainer.py", line 133, in __init__
    self.trainset, self.testset = self.get_dataset()
  File "C:\Users\MH_3002\anaconda3\envs\multitask-floorplan-3.10\lib\site-packages\ultralytics\engine\trainer.py", line 561, in get_dataset
    raise RuntimeError(emojis(f"Dataset '{clean_url(self.args.data)}' error ❌ {e}")) from e

Just noticed that YOLO rebuild the train() method and it doesn’t have the bool param. Thus, when a function call train(mode=False) it will call train() anyway.

But I still don’t have idea how to resolve it.

@shawn it might help to know your exact goal here. With what you shared, I feel like you could accomplish the same thing using:

from ultralytics import YOLO

model = YOLO("path/to/model.pt")
result = model.predict("path/to/source", save_crop=True)

or

yolo predict \
    model=path/to/model.pt \
    source=path/to/source \
    save_crop

You can also accomplish the same using an exported ONNX model.