refactor Accelerator with Fabric (#192)

* refactor Accelerator with Fabric

* update

* fix

* remove flash

* fixes

* update

* update

* update

Author: aniketmaurya

examples/src/models/hello_world.py

@@ -15,6 +15,7 @@
 import torchvision
 from timm import create_model
 from torch.utils.data import DataLoader
+from torchmetrics.classification import MulticlassAccuracy
 from torchvision import transforms as T
 
 from gradsflow import AutoDataset, Model
@@ -55,5 +56,5 @@
 
     model = Model(cnn)
 
-    model.compile("crossentropyloss", "adam", metrics=["accuracy"])
+    model.compile("crossentropyloss", "adam", metrics=[MulticlassAccuracy(autodataset.num_classes)])
     model.fit(autodataset, max_epochs=10, steps_per_epoch=10, callbacks=cbs)

gradsflow/autotasks/engine/backend.py

@@ -25,12 +25,12 @@
 from gradsflow.utility.imports import is_installed
 
 if typing.TYPE_CHECKING:
-    import pytorch_lightning as pl
+    import lightning as L
 
-if is_installed("pytorch_lightning"):
+if is_installed("lightning-flash"):
     from flash import Task
     from flash import Trainer as FlashTrainer
-    from pytorch_lightning import Trainer as PLTrainer
+    from lightning import Trainer as PLTrainer
 else:
     FlashTrainer = None
     PLTrainer = None
@@ -40,10 +40,10 @@
 
 class BackendType(Enum):
     # Remove torch
-    pl = "pl"
+    lightning = "lightning"
     gf = "gf"
     torch = "gf"
-    default = "pl"
+    default = "lightning"
 
 
 class Backend:
@@ -90,7 +90,7 @@ def _lightning_objective(
 
         trainer_cls = FlashTrainer if isinstance(model, Task) else PLTrainer
 
-        trainer: "pl.Trainer" = trainer_cls(
+        trainer: "L.Trainer" = trainer_cls(
             logger=True,
             accelerator="auto",
             devices="auto",
@@ -122,7 +122,7 @@ def optimization_objective(
             trainer_config dict: configurations passed directly to Lightning Trainer.
             gpu Optional[float]: GPU per trial
         """
-        if self.backend_type == BackendType.pl.value:
+        if self.backend_type == BackendType.lightning.value:
             return self._lightning_objective(config, trainer_config=trainer_config, gpu=gpu, finetune=finetune)
 
         if self.backend_type in (BackendType.gf.value,):

gradsflow/data/autodata.py

@@ -15,7 +15,7 @@
 import warnings
 from typing import Callable, Dict, Optional, Union
 
-from accelerate import Accelerator
+from lightning.fabric import Fabric
 from torch.utils.data import DataLoader, Dataset
 
 from gradsflow.data.base import BaseAutoDataset
@@ -131,13 +131,13 @@ def device_setup_status(self, value: bool = True):
         logger.debug("setting device setup=True")
         self.meta["device_setup_status"] = value
 
-    def prepare_data(self, accelerator: Accelerator) -> None:
+    def setup_data(self, accelerator: Fabric) -> None:
         if accelerator is None:
             warnings.warn("Accelerator is None, skipped data preparation!")
             return
-        self._train_dataloader = accelerator.prepare_data_loader(self._train_dataloader)
+        self._train_dataloader = accelerator.setup_dataloaders(self._train_dataloader)
         if self._val_dataloader:
-            self._val_dataloader = accelerator.prepare_data_loader(self._val_dataloader)
+            self._val_dataloader = accelerator.setup_dataloaders(self._val_dataloader)
         self.device_setup_status = True
         self.device = accelerator.device
 

gradsflow/models/base.py

@@ -11,13 +11,14 @@
 #  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.
+import logging
 import os
 from dataclasses import dataclass
 from typing import Any, Callable, List, Optional, Union
 
 import smart_open
 import torch
-from accelerate import Accelerator
+from lightning.fabric import Fabric
 from torch import nn
 
 from gradsflow.models.tracker import Tracker
@@ -31,7 +32,7 @@
 class Base:
     TEST = os.environ.get("GF_CI", "false").lower() == "true"
 
-    learner: Union[nn.Module, Any]
+    _learner: Union[nn.Module, Any]
     optimizer: torch.optim.Optimizer = None
     loss: Callable = None
     _compiled: bool = False
@@ -43,6 +44,14 @@ def __init__(self):
     def __call__(self, x):
         return self.forward(x)
 
+    @property
+    def learner(self) -> Union[nn.Module, Any]:
+        return self._learner
+
+    @learner.setter
+    def learner(self, learner):
+        self._learner = learner
+
     @staticmethod
     def _get_loss(loss: Union[str, Callable], loss_config: dict) -> Optional[Callable]:
         loss_fn = None
@@ -101,43 +110,24 @@ class BaseModel(Base):
     def __init__(
         self,
         learner: Union[nn.Module, Any],
-        device: Optional[str] = None,
-        use_accelerate: bool = True,
+        device: Optional[str] = "auto",
+        use_accelerator: bool = True,
         accelerator_config: dict = None,
     ):
         self.accelerator = None
         super().__init__()
-        self._set_accelerator(device, use_accelerate, accelerator_config)
-        self.learner = self.prepare_model(learner)
+        self._set_accelerator(device, use_accelerator, accelerator_config)
+        self._learner = learner
 
     def _set_accelerator(self, device: Optional[str], use_accelerate: bool, accelerator_config: dict):
         if use_accelerate:
-            self.accelerator = Accelerator(cpu=(device == "cpu"), **accelerator_config)
+            self.accelerator = Fabric(accelerator=device, **accelerator_config)
             self.device = self.accelerator.device
         else:
             self.device = device or default_device()
 
-    def prepare_model(self, learner: Union[nn.Module, List[nn.Module]]):
-        """Inplace ops for preparing model via HF Accelerator. Automatically sends to device."""
-        if not self.accelerator:
-            learner = learner.to(self.device)
-            return learner
-        if isinstance(learner, (list, tuple)):
-            self.learner = list(map(self.accelerator.prepare_model, learner))
-        elif isinstance(learner, nn.Module):
-            self.learner = self.accelerator.prepare_model(learner)
-        else:
-            raise NotImplementedError(
-                f"prepare_model is not implemented for model of type {type(learner)}! Please implement prepare_model "
-                f"or raise an issue."
-            )
-
-        return self.learner
-
-    def prepare_optimizer(self, optimizer) -> torch.optim.Optimizer:
-        if not self.accelerator:
-            return optimizer
-        return self.accelerator.prepare_optimizer(optimizer)
+    def setup(self, learner: Union[nn.Module, List[nn.Module]], *optimizers):
+        return self.accelerator.setup(learner, *optimizers)
 
     def backward(self, loss: torch.Tensor):
         """model.backward(loss)"""

gradsflow/models/model.py

@@ -46,7 +46,7 @@ class Model(BaseModel, DataMixin):
 
     Args:
         learner: Trainable model
-        accelerator_config: HuggingFace Accelerator config
+        accelerator_config: Accelerator config
     """
 
     TEST = os.environ.get("GF_CI", "false").lower() == "true"
@@ -56,14 +56,14 @@ def __init__(
         self,
         learner: Union[nn.Module, Any],
         device: Optional[str] = None,
-        use_accelerate: bool = True,
+        use_accelerator: bool = True,
         accelerator_config: dict = None,
     ):
         accelerator_config = accelerator_config or {}
         super().__init__(
             learner=learner,
             device=device,
-            use_accelerate=use_accelerate,
+            use_accelerator=use_accelerator,
             accelerator_config=accelerator_config,
         )
 
@@ -119,9 +119,9 @@ def compile(
         if optimizer:
             optimizer_fn = self._get_optimizer(optimizer)
             optimizer = optimizer_fn(self.learner.parameters(), lr=learning_rate, **optimizer_config)
-            self.optimizer = self.prepare_optimizer(optimizer)
         if loss:
             self.loss = self._get_loss(loss, loss_config)
+        self.learner, self.optimizer = self.setup(self._learner, optimizer)
         self.metrics.compile_metrics(*listify(metrics))
         self._compiled = True
 
@@ -244,7 +244,7 @@ def fit(
         """
         self.assert_compiled()
         self.autodataset = autodataset
-        self.autodataset.prepare_data(self.accelerator)
+        self.autodataset.setup_data(self.accelerator)
 
         if not resume:
             self.tracker.reset()

gradsflow/utility/common.py

@@ -13,6 +13,7 @@
 #  limitations under the License.
 import dataclasses
 import inspect
+import logging
 import os
 import re
 import sys
@@ -120,7 +121,7 @@ def to_item(data: Any) -> Union[int, float, str, np.ndarray, Dict]:
             data = data.detach()
         data = data.cpu().numpy()
 
-    warnings.warn("to_item didn't convert any value.")
+    logging.info("to_item didn't convert any value.")
     return data
 
 

pyproject.toml

@@ -11,3 +11,7 @@ profile = "black"
 
 [tool.black]
 line_length = 120
+
+
+[tool.pytest.ini_options]
+norecursedirs = ["tests/autotasks", "tests/tuner"]

setup.cfg

@@ -29,16 +29,16 @@ python_requires = >=3.8
 install_requires =
     torch >=1.13.1
     torchvision
-    ray[default,tune] >=1.8.0
+    ray[default,tune] >=2.2.0
     timm>=0.6.12
     rich>=13.3.1
-    accelerate >=0.5.0
     smart_open >=5.1,<=5.2.1
     torchmetrics >=0.11.1
+    lightning >=1.9.2
 
 [options.extras_require]
-dev = lightning-flash[image,text] >=0.5.1; codecarbon >=1.2.0; comet_ml; wandb; tensorboard
-test = pytest; coverage; pytest-sugar
+dev = codecarbon >=1.2.0; wandb; tensorboard
+test = pytest; coverage; pytest-sugar; pytest-randomly
 
 [options.packages.find] #optional
 exclude=tests, docs, examples

tests/__main__.py

@@ -12,15 +12,17 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import urllib.request
+import zipfile
 from pathlib import Path
 
-from flash.core.data.utils import download_data
+cwd = Path.cwd()
+(Path.cwd() / "data").mkdir(exist_ok=True)
 
-cwd = str(Path.cwd())
-
-download_data("https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip", f"{cwd}/data")
-
-download_data(
+urllib.request.urlretrieve(
     "https://github.com/gradsflow/test-data/archive/refs/tags/cat-dog-v0.zip",
-    f"{cwd}/data",
+    f"{cwd}/data/test-cat-dog-v0.zip",
 )
+
+with zipfile.ZipFile(f"{cwd}/data/test-cat-dog-v0.zip", "r") as zip_ref:
+    zip_ref.extractall(f"{cwd}/data/")

tests/autotasks/test_core_automodel.py

@@ -55,7 +55,7 @@ def test_objective(mock_pl_trainer, mock_fl_trainer):
     model = AutoModel(
         datamodule,
         optimization_metric=optimization_metric,
-        backend=BackendType.pl.value,
+        backend=BackendType.lightning.value,
     )
 
     model.backend.model_builder = MagicMock()