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’, …
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.