feat: allow custom optimizer

src/pytorch_tabular/config/config.py

@@ -655,7 +655,8 @@ class OptimizerConfig:
     """Optimizer and Learning Rate Scheduler configuration.
     Args:
         optimizer (str): Any of the standard optimizers from
-                [torch.optim](https://pytorch.org/docs/stable/optim.html#algorithms).
+                [torch.optim](https://pytorch.org/docs/stable/optim.html#algorithms) or provide full python path,
+                for example "torch_optimizer.RAdam".
 
         optimizer_params (Dict): The parameters for the optimizer. If left blank, will use default
                 parameters.
@@ -675,7 +676,8 @@ class OptimizerConfig:
         default="Adam",
         metadata={
             "help": "Any of the standard optimizers from"
-            " [torch.optim](https://pytorch.org/docs/stable/optim.html#algorithms)."
+            " [torch.optim](https://pytorch.org/docs/stable/optim.html#algorithms) or provide full python path,"
+            " for example 'torch_optimizer.RAdam'."
         },
     )
     optimizer_params: Dict = field(

src/pytorch_tabular/ssl_models/base_model.py

@@ -2,6 +2,7 @@
 # Author: Manu Joseph <[email protected]>
 # For license information, see LICENSE.TXT
 """SSL Base Model."""
+import importlib
 import warnings
 from abc import ABCMeta, abstractmethod
 from typing import Dict, Optional
@@ -172,7 +173,12 @@ def configure_optimizers(self):
         if self.custom_optimizer is None:
             # Loading from the config
             try:
-                self._optimizer = getattr(torch.optim, self.hparams.optimizer)
+                if "." in self.hparams.optimizer:
+                    py_path, cls_name = self.hparams.optimizer.rsplit(".", 1)
+                    module = importlib.import_module(py_path)
+                    self._optimizer = getattr(module, cls_name)
+                else:
+                    self._optimizer = getattr(torch.optim, self.hparams.optimizer)
                 opt = self._optimizer(
                     self.parameters(),
                     lr=self.hparams.learning_rate,

src/pytorch_tabular/tabular_model.py

@@ -510,10 +510,10 @@ def prepare_model(
                 The length of the list should be equal to the number of metrics. Defaults to None.
 
             optimizer (Optional[torch.optim.Optimizer], optional):
-                Custom optimizers which are a drop in replacements for standard PyToch optimizers.
+                Custom optimizers which are a drop in replacements for standard PyTorch optimizers.
                 This should be the Class and not the initialized object
 
-            optimizer_params (Optional[Dict], optional): The parmeters to initialize the custom optimizer.
+            optimizer_params (Optional[Dict], optional): The parameters to initialize the custom optimizer.
 
         Returns:
             BaseModel: The prepared model
@@ -631,9 +631,9 @@ def fit(
 
             optimizer (Optional[torch.optim.Optimizer], optional):
                 Custom optimizers which are a drop in replacements for
-                standard PyToch optimizers. This should be the Class and not the initialized object
+                standard PyTorch optimizers. This should be the Class and not the initialized object
 
-            optimizer_params (Optional[Dict], optional): The parmeters to initialize the custom optimizer.
+            optimizer_params (Optional[Dict], optional): The parameters to initialize the custom optimizer.
 
             train_sampler (Optional[torch.utils.data.Sampler], optional):
                 Custom PyTorch batch samplers which will be passed
@@ -717,9 +717,9 @@ def pretrain(
                 Defaults to None.
 
             optimizer (Optional[torch.optim.Optimizer], optional): Custom optimizers which are a drop in replacements
-                for standard PyToch optimizers. This should be the Class and not the initialized object
+                for standard PyTorch optimizers. This should be the Class and not the initialized object
 
-            optimizer_params (Optional[Dict], optional): The parmeters to initialize the custom optimizer.
+            optimizer_params (Optional[Dict], optional): The parameters to initialize the custom optimizer.
 
             max_epochs (Optional[int]): Overwrite maximum number of epochs to be run. Defaults to None.
 
@@ -731,7 +731,7 @@ def pretrain(
                 Defaults to None.
 
             datamodule (Optional[TabularDatamodule], optional): The datamodule. If provided, will ignore the rest of the
-                parameters like train, test etc and use the datamodule. Defaults to None.
+                parameters like train, test etc. and use the datamodule. Defaults to None.
 
         Returns:
             pl.Trainer: The PyTorch Lightning Trainer instance
@@ -808,7 +808,7 @@ def create_finetune_model(
 
             loss (Optional[torch.nn.Module], optional):
                 If provided, will be used as the loss function for the fine-tuning.
-                By Default it is MSELoss for regression and CrossEntropyLoss for classification.
+                By default, it is MSELoss for regression and CrossEntropyLoss for classification.
 
             metrics (Optional[List[Callable]], optional): List of metrics (either callables or str) to be used for the
                 fine-tuning stage. If str, it should be one of the functional metrics implemented in