PoC: Revamp optimizer and scheduler experience using registries

flash/audio/speech_recognition/model.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import os
 import warnings
-from typing import Any, Dict, Mapping, Optional, Type, Union
+from typing import Any, Callable, Dict, Mapping, Optional, Type, Union
 
 import torch
 import torch.nn as nn
@@ -54,8 +54,8 @@ class SpeechRecognition(Task):
     def __init__(
         self,
         backbone: str = "facebook/wav2vec2-base-960h",
-        optimizer: Type[torch.optim.Optimizer] = torch.optim.Adam,
-        optimizer_kwargs: Optional[Dict[str, Any]] = None,
+        optimizer: Union[Callable[..., torch.optim.Optimizer], str] = "Adam",
+        # optimizer_kwargs: Optional[Dict[str, Any]] = None,
         scheduler: Optional[Union[Type[_LRScheduler], str, _LRScheduler]] = None,
         scheduler_kwargs: Optional[Dict[str, Any]] = None,
         learning_rate: float = 1e-5,
@@ -71,7 +71,7 @@ def __init__(
         super().__init__(
             model=model,
             optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            # optimizer_kwargs=optimizer_kwargs,
             scheduler=scheduler,
             scheduler_kwargs=scheduler_kwargs,
             learning_rate=learning_rate,

flash/core/model.py

@@ -46,9 +46,10 @@
     SerializerMapping,
 )
 from flash.core.data.properties import ProcessState
+from flash.core.optimizers import _OPTIMIZERS_REGISTRY, _SCHEDULERS_REGISTRY
 from flash.core.registry import FlashRegistry
-from flash.core.schedulers import _SCHEDULERS_REGISTRY
 from flash.core.serve import Composition
+from flash.core.utilities import providers
 from flash.core.utilities.apply_func import get_callable_dict
 from flash.core.utilities.imports import requires
 
@@ -303,6 +304,7 @@ class Task(DatasetProcessor, ModuleWrapperBase, LightningModule, metaclass=Check
         postprocess: :class:`~flash.core.data.process.Postprocess` to use as the default for this task.
     """
 
+    optimizers: FlashRegistry = _OPTIMIZERS_REGISTRY
     schedulers: FlashRegistry = _SCHEDULERS_REGISTRY
 
     required_extras: Optional[Union[str, List[str]]] = None
@@ -311,12 +313,12 @@ def __init__(
         self,
         model: Optional[nn.Module] = None,
         loss_fn: Optional[Union[Callable, Mapping, Sequence]] = None,
-        optimizer: Union[Type[torch.optim.Optimizer], torch.optim.Optimizer] = torch.optim.Adam,
-        optimizer_kwargs: Optional[Dict[str, Any]] = None,
-        scheduler: Optional[Union[Type[_LRScheduler], str, _LRScheduler]] = None,
-        scheduler_kwargs: Optional[Dict[str, Any]] = None,
-        metrics: Union[torchmetrics.Metric, Mapping, Sequence, None] = None,
         learning_rate: float = 5e-5,
+        optimizer: Union[Callable[..., torch.optim.Optimizer], str] = "Adam",
+        # optimizer_kwargs: Optional[Dict[str, Any]] = None,
+        scheduler: Optional[Union[str, Callable, Tuple[str, Tuple[Any, ...]]]] = None,
+        # scheduler_kwargs: Optional[Dict[str, Any]] = None,
+        metrics: Union[torchmetrics.Metric, Mapping, Sequence, None] = None,
         deserializer: Optional[Union[Deserializer, Mapping[str, Deserializer]]] = None,
         preprocess: Optional[Preprocess] = None,
         postprocess: Optional[Postprocess] = None,
@@ -328,8 +330,10 @@ def __init__(
         self.loss_fn = {} if loss_fn is None else get_callable_dict(loss_fn)
         self.optimizer = optimizer
         self.scheduler = scheduler
-        self.optimizer_kwargs = optimizer_kwargs or {}
-        self.scheduler_kwargs = scheduler_kwargs or {}
+        if isinstance(self.scheduler, Tuple):
+            assert isinstance(self.scheduler[0], str)
+        # self.optimizer_kwargs: Dict[str, Any] = optimizer_kwargs or {}
+        # self.scheduler_kwargs: Dict[str, Any] = scheduler_kwargs or {}
 
         self.train_metrics = nn.ModuleDict({} if metrics is None else get_callable_dict(metrics))
         self.val_metrics = nn.ModuleDict({} if metrics is None else get_callable_dict(deepcopy(metrics)))
@@ -474,11 +478,15 @@ def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> A
         return self(batch)
 
     def configure_optimizers(self) -> Union[Optimizer, Tuple[List[Optimizer], List[_LRScheduler]]]:
-        optimizer = self.optimizer
-        if not isinstance(self.optimizer, Optimizer):
-            self.optimizer_kwargs["lr"] = self.learning_rate
-            optimizer = optimizer(filter(lambda p: p.requires_grad, self.parameters()), **self.optimizer_kwargs)
-        if self.scheduler:
+        if isinstance(self.optimizer, str):
+            self.optimizer = self.optimizers.get(self.optimizer)
+
+        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
+        optimizer: Optimizer = self.optimizer(model_parameters, lr=self.learning_rate)
+        # if not isinstance(self.optimizer, Optimizer):
+        #     self.optimizer_kwargs["lr"] = self.learning_rate
+        #     optimizer = optimizer(filter(lambda p: p.requires_grad, self.parameters()), **self.optimizer_kwargs)
+        if self.scheduler is not None:
             return [optimizer], [self._instantiate_scheduler(optimizer)]
         return optimizer
 
@@ -817,23 +825,50 @@ def _compute_warmup(num_training_steps: int, num_warmup_steps: Union[int, float]
         return round(num_warmup_steps)
 
     def _instantiate_scheduler(self, optimizer: Optimizer) -> _LRScheduler:
-        scheduler = self.scheduler
-        if isinstance(scheduler, _LRScheduler):
-            return scheduler
-        if isinstance(scheduler, str):
-            scheduler_fn = self.schedulers.get(self.scheduler)
-            num_training_steps: int = self.get_num_training_steps()
-            num_warmup_steps: int = self._compute_warmup(
-                num_training_steps=num_training_steps,
-                num_warmup_steps=self.scheduler_kwargs.get("num_warmup_steps"),
-            )
-            return scheduler_fn(optimizer, num_warmup_steps, num_training_steps)
-        if issubclass(scheduler, _LRScheduler):
-            return scheduler(optimizer, **self.scheduler_kwargs)
-        raise MisconfigurationException(
-            "scheduler can be a scheduler, a scheduler type with `scheduler_kwargs` "
-            f"or a built-in scheduler in {self.available_schedulers()}"
-        )
+        if isinstance(self.scheduler, Callable):
+            return self.scheduler(optimizer)
+
+        if isinstance(self.scheduler, str):
+            return self.schedulers.get(self.scheduler)(optimizer)
+
+        if not isinstance(self.scheduler, Tuple):
+            raise TypeError("")  # Add message
+
+        # By default it is Tuple[str, Tuple[Any, ...]] type now.
+        scheduler_key = self.scheduler[0]
+        scheduler_args = self.scheduler[1]
+
+        scheduler = self.schedulers.get(scheduler_key, with_metadata=True)
+        scheduler_fn = scheduler["fn"]
+        scheduler_metadata = scheduler["metadata"]
+
+        if "providers" in scheduler_metadata.keys():
+            if scheduler_metadata["providers"] == providers._HUGGINGFACE:
+                num_training_steps: int = self.get_num_training_steps()
+                num_warmup_steps: int = self._compute_warmup(
+                    num_training_steps=num_training_steps,
+                    num_warmup_steps=scheduler_args[0],  # num_warmup_steps is the first arg in all schedulers
+                )
+                if scheduler["name"] == "constant_schedule_with_warmup":
+                    scheduler_args = (num_warmup_steps, *scheduler_args[1:])
+                else:
+                    scheduler_args = (num_warmup_steps, num_training_steps, *scheduler_args[1:])
+
+            # if isinstance(scheduler, str):
+            #     scheduler_fn = self.schedulers.get(self.scheduler)
+            #     num_training_steps: int = self.get_num_training_steps()
+            #     num_warmup_steps: int = self._compute_warmup(
+            #         num_training_steps=num_training_steps,
+            #         num_warmup_steps=self.scheduler_kwargs.get("num_warmup_steps"),
+            #     )
+            #     return scheduler_fn(optimizer, num_warmup_steps, num_training_steps)
+            # # if issubclass(scheduler, _LRScheduler):
+            # #     return scheduler(optimizer, **self.scheduler_kwargs)
+            # raise MisconfigurationException(
+            #     "scheduler can be a scheduler, a scheduler type with `scheduler_kwargs` "
+            #     f"or a built-in scheduler in {self.available_schedulers()}"
+            # )
+        return scheduler_fn(optimizer, *scheduler_args)
 
     def _load_from_state_dict(
         self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs

flash/core/optimizers/__init__.py

@@ -1,3 +1,5 @@
 from flash.core.optimizers.lamb import LAMB  # noqa: F401
 from flash.core.optimizers.lars import LARS  # noqa: F401
 from flash.core.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR  # noqa: F401
+from flash.core.optimizers.optimizers import _OPTIMIZERS_REGISTRY  # noqa: F401
+from flash.core.optimizers.schedulers import _SCHEDULERS_REGISTRY  # noqa: F401

flash/core/optimizers/optimizers.py

@@ -0,0 +1,12 @@
+from typing import Callable, List
+
+from torch import optim
+
+from flash.core.registry import FlashRegistry
+
+_OPTIMIZERS_REGISTRY = FlashRegistry("optimizer")
+
+_optimizers: List[Callable] = [getattr(optim, n) for n in dir(optim) if ("_" not in n)]
+
+for fn in _optimizers:
+    _OPTIMIZERS_REGISTRY(fn, name=fn.__name__)

flash/core/optimizers/schedulers.py

@@ -0,0 +1,33 @@
+import inspect
+from typing import Callable, List
+
+from torch.optim import lr_scheduler
+from torch.optim.lr_scheduler import _LRScheduler
+
+from flash.core.registry import FlashRegistry
+from flash.core.utilities.imports import _TRANSFORMERS_AVAILABLE
+from flash.core.utilities.providers import _HUGGINGFACE
+
+_SCHEDULERS_REGISTRY = FlashRegistry("scheduler")
+
+schedulers: List[_LRScheduler] = []
+for n in dir(lr_scheduler):
+    sched = getattr(lr_scheduler, n)
+
+    if inspect.isclass(sched) and sched != _LRScheduler and issubclass(sched, _LRScheduler):
+        schedulers.append(sched)
+
+
+for scheduler in schedulers:
+    _SCHEDULERS_REGISTRY(scheduler, name=scheduler.__name__)
+
+if _TRANSFORMERS_AVAILABLE:
+    from transformers import optimization
+
+    functions: List[Callable] = []
+    for n in dir(optimization):
+        if "get_" in n and n != "get_scheduler":
+            functions.append(getattr(optimization, n))
+
+    for fn in functions:
+        _SCHEDULERS_REGISTRY(fn, name=fn.__name__[4:], providers=_HUGGINGFACE)

flash/graph/classification/model.py

@@ -109,8 +109,8 @@ def __init__(
         num_classes: int,
         hidden_channels: Union[List[int], int] = 512,
         loss_fn: Callable = F.cross_entropy,
-        optimizer: Type[torch.optim.Optimizer] = torch.optim.Adam,
-        optimizer_kwargs: Optional[Dict[str, Any]] = None,
+        optimizer: Union[Callable[..., torch.optim.Optimizer], str] = "Adam",
+        # optimizer_kwargs: Optional[Dict[str, Any]] = None,
         scheduler: Optional[Union[Type[_LRScheduler], str, _LRScheduler]] = None,
         scheduler_kwargs: Optional[Dict[str, Any]] = None,
         metrics: Union[Callable, Mapping, Sequence, None] = None,
@@ -132,7 +132,7 @@ def __init__(
             model=model,
             loss_fn=loss_fn,
             optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            # optimizer_kwargs=optimizer_kwargs,
             scheduler=scheduler,
             scheduler_kwargs=scheduler_kwargs,
             metrics=metrics,

flash/image/classification/model.py

@@ -81,8 +81,8 @@ def __init__(
         head: Optional[Union[FunctionType, nn.Module]] = None,
         pretrained: Union[bool, str] = True,
         loss_fn: Optional[Callable] = None,
-        optimizer: Union[Type[torch.optim.Optimizer], torch.optim.Optimizer] = torch.optim.Adam,
-        optimizer_kwargs: Optional[Dict[str, Any]] = None,
+        optimizer: Union[Callable[..., torch.optim.Optimizer], str] = "Adam",
+        # optimizer_kwargs: Optional[Dict[str, Any]] = None,
         scheduler: Optional[Union[Type[_LRScheduler], str, _LRScheduler]] = None,
         scheduler_kwargs: Optional[Dict[str, Any]] = None,
         metrics: Union[Metric, Callable, Mapping, Sequence, None] = None,
@@ -138,7 +138,7 @@ def __init__(
             metrics=metrics,
             learning_rate=learning_rate,
             optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            # optimizer_kwargs=optimizer_kwargs,
             scheduler=scheduler,
             scheduler_kwargs=scheduler_kwargs,
             multi_label=multi_label,

flash/image/detection/model.py

@@ -11,10 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, List, Mapping, Optional, Type, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Type, Union
 
 import torch
-from torch.optim import Optimizer
 from torch.optim.lr_scheduler import _LRScheduler
 
 from flash.core.adapter import AdapterTask
@@ -61,8 +60,8 @@ def __init__(
         backbone: Optional[str] = "resnet18_fpn",
         head: Optional[str] = "retinanet",
         pretrained: bool = True,
-        optimizer: Type[Optimizer] = torch.optim.Adam,
-        optimizer_kwargs: Optional[Dict[str, Any]] = None,
+        optimizer: Union[Callable[..., torch.optim.Optimizer], str] = "Adam",
+        # optimizer_kwargs: Optional[Dict[str, Any]] = None,
         scheduler: Optional[Union[Type[_LRScheduler], str, _LRScheduler]] = None,
         scheduler_kwargs: Optional[Dict[str, Any]] = None,
         learning_rate: float = 5e-3,
@@ -85,7 +84,7 @@ def __init__(
             adapter,
             learning_rate=learning_rate,
             optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            # optimizer_kwargs=optimizer_kwargs,
             scheduler=scheduler,
             scheduler_kwargs=scheduler_kwargs,
             serializer=serializer or Preds(),

flash/image/embedding/model.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import warnings
-from typing import Any, Dict, List, Optional, Type, Union
+from typing import Any, Callable, Dict, List, Optional, Type, Union
 
 import torch
 from torch.optim.lr_scheduler import _LRScheduler
@@ -77,8 +77,8 @@ def __init__(
         pretraining_transform: str,
         backbone: str = "resnet",
         pretrained: bool = False,
-        optimizer: Type[torch.optim.Optimizer] = torch.optim.SGD,
-        optimizer_kwargs: Optional[Dict[str, Any]] = None,
+        optimizer: Union[Callable[..., torch.optim.Optimizer], str] = "SGD",
+        # optimizer_kwargs: Optional[Dict[str, Any]] = None,
         scheduler: Optional[Union[Type[_LRScheduler], str, _LRScheduler]] = None,
         scheduler_kwargs: Optional[Dict[str, Any]] = None,
         learning_rate: float = 1e-3,
@@ -110,7 +110,7 @@ def __init__(
         super().__init__(
             adapter=adapter,
             optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            # optimizer_kwargs=optimizer_kwargs,
             scheduler=scheduler,
             scheduler_kwargs=scheduler_kwargs,
             learning_rate=learning_rate,

flash/image/instance_segmentation/model.py

@@ -11,10 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, List, Mapping, Optional, Type, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Type, Union
 
 import torch
-from torch.optim import Optimizer
 from torch.optim.lr_scheduler import _LRScheduler
 
 from flash.core.adapter import AdapterTask
@@ -61,8 +60,8 @@ def __init__(
         backbone: Optional[str] = "resnet18_fpn",
         head: Optional[str] = "mask_rcnn",
         pretrained: bool = True,
-        optimizer: Type[Optimizer] = torch.optim.Adam,
-        optimizer_kwargs: Optional[Dict[str, Any]] = None,
+        optimizer: Union[Callable[..., torch.optim.Optimizer], str] = "Adam",
+        # optimizer_kwargs: Optional[Dict[str, Any]] = None,
         scheduler: Optional[Union[Type[_LRScheduler], str, _LRScheduler]] = None,
         scheduler_kwargs: Optional[Dict[str, Any]] = None,
         learning_rate: float = 5e-4,
@@ -85,7 +84,7 @@ def __init__(
             adapter,
             learning_rate=learning_rate,
             optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            # optimizer_kwargs=optimizer_kwargs,
             scheduler=scheduler,
             scheduler_kwargs=scheduler_kwargs,
             serializer=serializer or Preds(),