Refactor optimizer and lr scheduler part (#3216)

* Refactor optimizer and lr schedulers * Update tests * Update recipe * Fix torch import problem at CLI installation Signed-off-by: Kim, Vinnam <[email protected]> * Fix Signed-off-by: Kim, Vinnam <[email protected]> * Fix test errors Signed-off-by: Kim, Vinnam <[email protected]> * Fix Signed-off-by: Kim, Vinnam <[email protected]> * Rollback how OVVisualPromptingModel get model names Signed-off-by: Kim, Vinnam <[email protected]> * Update src/otx/algo/anomaly/openvino_model.py Co-authored-by: Harim Kang <[email protected]> --------- Signed-off-by: Kim, Vinnam <[email protected]> Co-authored-by: Harim Kang <[email protected]>
openvinotoolkit · Mar 29, 2024 · 375e89e · 375e89e
1 parent 9a47ee8
commit 375e89e
Show file tree

Hide file tree

Showing 123 changed files with 1,626 additions and 1,340 deletions.
diff --git a/src/otx/algo/action_classification/movinet.py b/src/otx/algo/action_classification/movinet.py
@@ -12,6 +12,7 @@
 from otx.core.metrics.accuracy import MultiClassClsMetricCallable
 from otx.core.model.action_classification import MMActionCompatibleModel
 from otx.core.model.base import DefaultOptimizerCallable, DefaultSchedulerCallable
+from otx.core.schedulers import LRSchedulerListCallable
 
 if TYPE_CHECKING:
     from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
@@ -25,8 +26,8 @@ class MoViNet(MMActionCompatibleModel):
     def __init__(
         self,
         num_classes: int,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MultiClassClsMetricCallable,
         torch_compile: bool = False,
     ) -> None:

diff --git a/src/otx/algo/action_classification/x3d.py b/src/otx/algo/action_classification/x3d.py
@@ -11,6 +11,7 @@
 from otx.core.metrics.accuracy import MultiClassClsMetricCallable
 from otx.core.model.action_classification import MMActionCompatibleModel
 from otx.core.model.base import DefaultOptimizerCallable, DefaultSchedulerCallable
+from otx.core.schedulers import LRSchedulerListCallable
 
 if TYPE_CHECKING:
     from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
@@ -24,8 +25,8 @@ class X3D(MMActionCompatibleModel):
     def __init__(
         self,
         num_classes: int,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MultiClassClsMetricCallable,
         torch_compile: bool = False,
     ) -> None:

diff --git a/src/otx/algo/action_detection/x3d_fastrcnn.py b/src/otx/algo/action_detection/x3d_fastrcnn.py
@@ -11,6 +11,7 @@
 from otx.core.metrics.mean_ap import MeanAPCallable
 from otx.core.model.action_detection import MMActionCompatibleModel
 from otx.core.model.base import DefaultOptimizerCallable, DefaultSchedulerCallable
+from otx.core.schedulers import LRSchedulerListCallable
 
 if TYPE_CHECKING:
     from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
@@ -25,8 +26,8 @@ def __init__(
         self,
         num_classes: int,
         topk: int | tuple[int],
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MeanAPCallable,
         torch_compile: bool = False,
     ) -> None:

diff --git a/src/otx/algo/anomaly/openvino_model.py b/src/otx/algo/anomaly/openvino_model.py
@@ -2,6 +2,7 @@
 
 All anomaly models use the same AnomalyDetection model from ModelAPI.
 """
+
 # TODO(someone): Revisit mypy errors after OTXLitModule deprecation and anomaly refactoring
 # mypy: ignore-errors
 
@@ -12,17 +13,16 @@
 
 from typing import TYPE_CHECKING, Any
 
-from lightning.pytorch import LightningModule
-
+from otx.core.metrics.types import MetricCallable, NullMetricCallable
 from otx.core.model.anomaly import AnomalyModelInputs
-from otx.core.model.base import OTXModel, OVModel
+from otx.core.model.base import OVModel
 
 if TYPE_CHECKING:
     from openvino.model_api.models import Model
     from openvino.model_api.models.anomaly import AnomalyResult
 
 
-class AnomalyOpenVINO(OVModel, OTXModel, LightningModule):
+class AnomalyOpenVINO(OVModel):
     """Anomaly OpenVINO model."""
 
     # [TODO](ashwinvaidya17): Remove LightningModule once OTXModel is updated to use LightningModule.
@@ -36,6 +36,7 @@ def __init__(
         use_throughput_mode: bool = True,
         model_api_configuration: dict[str, Any] | None = None,
         num_classes: int = 2,
+        metric: MetricCallable = NullMetricCallable,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -46,6 +47,7 @@ def __init__(
             max_num_requests=max_num_requests,
             use_throughput_mode=use_throughput_mode,
             model_api_configuration=model_api_configuration,
+            metric=metric,
         )
 
     def _create_model(self) -> Model:

diff --git a/src/otx/algo/classification/deit_tiny.py b/src/otx/algo/classification/deit_tiny.py
@@ -21,6 +21,7 @@
     MMPretrainMulticlassClsModel,
     MMPretrainMultilabelClsModel,
 )
+from otx.core.schedulers import LRSchedulerListCallable
 from otx.core.types.label import HLabelInfo
 
 if TYPE_CHECKING:
@@ -150,8 +151,8 @@ class DeitTinyForHLabelCls(ExplainableDeit, MMPretrainHlabelClsModel):
     def __init__(
         self,
         hlabel_info: HLabelInfo,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = HLabelClsMetricCallble,
         torch_compile: bool = False,
     ) -> None:
@@ -177,8 +178,8 @@ class DeitTinyForMulticlassCls(ExplainableDeit, MMPretrainMulticlassClsModel):
     def __init__(
         self,
         num_classes: int,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MultiClassClsMetricCallable,
         torch_compile: bool = False,
     ) -> None:
@@ -203,8 +204,8 @@ class DeitTinyForMultilabelCls(ExplainableDeit, MMPretrainMultilabelClsModel):
     def __init__(
         self,
         num_classes: int,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MultiLabelClsMetricCallable,
         torch_compile: bool = False,
     ) -> None:

diff --git a/src/otx/algo/classification/efficientnet_b0.py b/src/otx/algo/classification/efficientnet_b0.py
@@ -15,6 +15,7 @@
     MMPretrainMulticlassClsModel,
     MMPretrainMultilabelClsModel,
 )
+from otx.core.schedulers import LRSchedulerListCallable
 from otx.core.types.label import HLabelInfo
 
 if TYPE_CHECKING:
@@ -29,8 +30,8 @@ class EfficientNetB0ForHLabelCls(MMPretrainHlabelClsModel):
     def __init__(
         self,
         hlabel_info: HLabelInfo,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = HLabelClsMetricCallble,
         torch_compile: bool = False,
     ) -> None:
@@ -57,8 +58,8 @@ def __init__(
         self,
         num_classes: int,
         light: bool = False,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MultiClassClsMetricCallable,
         torch_compile: bool = False,
     ) -> None:
@@ -84,8 +85,8 @@ class EfficientNetB0ForMultilabelCls(MMPretrainMultilabelClsModel):
     def __init__(
         self,
         num_classes: int,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MultiLabelClsMetricCallable,
         torch_compile: bool = False,
     ) -> None:

diff --git a/src/otx/algo/classification/efficientnet_v2.py b/src/otx/algo/classification/efficientnet_v2.py
@@ -15,6 +15,7 @@
     MMPretrainMulticlassClsModel,
     MMPretrainMultilabelClsModel,
 )
+from otx.core.schedulers import LRSchedulerListCallable
 from otx.core.types.label import HLabelInfo
 
 if TYPE_CHECKING:
@@ -29,8 +30,8 @@ class EfficientNetV2ForHLabelCls(MMPretrainHlabelClsModel):
     def __init__(
         self,
         hlabel_info: HLabelInfo,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = HLabelClsMetricCallble,
         torch_compile: bool = False,
     ) -> None:
@@ -57,8 +58,8 @@ def __init__(
         self,
         num_classes: int,
         light: bool = False,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MultiClassClsMetricCallable,
         torch_compile: bool = False,
     ) -> None:
@@ -84,8 +85,8 @@ class EfficientNetV2ForMultilabelCls(MMPretrainMultilabelClsModel):
     def __init__(
         self,
         num_classes: int,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MultiLabelClsMetricCallable,
         torch_compile: bool = False,
     ) -> None:

diff --git a/src/otx/algo/classification/mmconfigs/multiclass_classification/dino_v2.yaml b/src/otx/algo/classification/mmconfigs/multiclass_classification/dino_v2.yaml
@@ -0,0 +1,6 @@
+backbone:
+  name: dinov2_vits14_reg
+  frozen: false
+head:
+  in_channels: 384
+  num_classes: 1000
diff --git a/src/otx/algo/classification/mobilenet_v3_large.py b/src/otx/algo/classification/mobilenet_v3_large.py
@@ -15,6 +15,7 @@
     MMPretrainMulticlassClsModel,
     MMPretrainMultilabelClsModel,
 )
+from otx.core.schedulers import LRSchedulerListCallable
 from otx.core.types.label import HLabelInfo
 
 if TYPE_CHECKING:
@@ -29,8 +30,8 @@ class MobileNetV3ForHLabelCls(MMPretrainHlabelClsModel):
     def __init__(
         self,
         hlabel_info: HLabelInfo,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = HLabelClsMetricCallble,
         torch_compile: bool = False,
     ) -> None:
@@ -64,8 +65,8 @@ def __init__(
         self,
         num_classes: int,
         light: bool = False,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MultiClassClsMetricCallable,
         torch_compile: bool = False,
     ) -> None:
@@ -98,8 +99,8 @@ class MobileNetV3ForMultilabelCls(MMPretrainMultilabelClsModel):
     def __init__(
         self,
         num_classes: int,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MultiLabelClsMetricCallable,
         torch_compile: bool = False,
     ) -> None:

diff --git a/src/otx/algo/classification/otx_dino_v2.py b/src/otx/algo/classification/otx_dino_v2.py
@@ -10,6 +10,7 @@
 import torch
 from torch import nn
 
+from otx.algo.utils.mmconfig import read_mmconfig
 from otx.core.data.entity.base import OTXBatchLossEntity
 from otx.core.data.entity.classification import (
     MulticlassClsBatchDataEntity,
@@ -20,11 +21,11 @@
 from otx.core.metrics.accuracy import MultiClassClsMetricCallable
 from otx.core.model.base import DefaultOptimizerCallable, DefaultSchedulerCallable
 from otx.core.model.classification import OTXMulticlassClsModel
+from otx.core.schedulers import LRSchedulerListCallable
 from otx.core.utils.config import inplace_num_classes
 
 if TYPE_CHECKING:
     from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
-    from omegaconf import DictConfig
 
     from otx.core.metrics import MetricCallable
 
@@ -76,14 +77,18 @@ class DINOv2RegisterClassifier(OTXMulticlassClsModel):
     def __init__(
         self,
         num_classes: int,
-        config: DictConfig,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MultiClassClsMetricCallable,
         torch_compile: bool = False,
+        freeze_backbone: bool = False,
     ) -> None:
+        config = read_mmconfig(model_name="dino_v2", subdir_name="multiclass_classification")
         config = inplace_num_classes(cfg=config, num_classes=num_classes)
+        config.backbone.frozen = freeze_backbone
+
         self.config = config
+
         super().__init__(
             num_classes=num_classes,
             optimizer=optimizer,

diff --git a/src/otx/algo/classification/torchvision_model.py b/src/otx/algo/classification/torchvision_model.py
@@ -21,6 +21,7 @@
 from otx.core.metrics.accuracy import MultiClassClsMetricCallable
 from otx.core.model.base import DefaultOptimizerCallable, DefaultSchedulerCallable
 from otx.core.model.classification import OTXMulticlassClsModel
+from otx.core.schedulers import LRSchedulerListCallable
 
 if TYPE_CHECKING:
     from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
@@ -182,8 +183,8 @@ def __init__(
         backbone: TVModelType,
         num_classes: int,
         loss_callable: Callable[[], nn.Module] = nn.CrossEntropyLoss,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MultiClassClsMetricCallable,
         torch_compile: bool = False,
         freeze_backbone: bool = False,

diff --git a/src/otx/algo/detection/atss.py b/src/otx/algo/detection/atss.py
@@ -12,6 +12,7 @@
 from otx.core.metrics.mean_ap import MeanAPCallable
 from otx.core.model.base import DefaultOptimizerCallable, DefaultSchedulerCallable
 from otx.core.model.detection import MMDetCompatibleModel
+from otx.core.schedulers import LRSchedulerListCallable
 
 if TYPE_CHECKING:
     from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
@@ -26,8 +27,8 @@ def __init__(
         self,
         num_classes: int,
         variant: Literal["mobilenetv2", "resnext101"],
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MeanAPCallable,
         torch_compile: bool = False,
     ) -> None:
@@ -67,8 +68,8 @@ class ATSSR50FPN(MMDetCompatibleModel):
     def __init__(
         self,
         num_classes: int,
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MeanAPCallable,
         torch_compile: bool = False,
     ) -> None:

diff --git a/src/otx/algo/detection/rtmdet.py b/src/otx/algo/detection/rtmdet.py
@@ -12,6 +12,7 @@
 from otx.core.metrics.mean_ap import MeanAPCallable
 from otx.core.model.base import DefaultOptimizerCallable, DefaultSchedulerCallable
 from otx.core.model.detection import MMDetCompatibleModel
+from otx.core.schedulers import LRSchedulerListCallable
 
 if TYPE_CHECKING:
     from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
@@ -26,8 +27,8 @@ def __init__(
         self,
         num_classes: int,
         variant: Literal["tiny"],
-        optimizer: list[OptimizerCallable] | OptimizerCallable = DefaultOptimizerCallable,
-        scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = DefaultSchedulerCallable,
+        optimizer: OptimizerCallable = DefaultOptimizerCallable,
+        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = MeanAPCallable,
         torch_compile: bool = False,
     ) -> None: