Test fabric (#193)

* refactor apis

* update

* fix

* update

* update

* fix test

Author: aniketmaurya

examples/src/models/hello_world.py

@@ -11,50 +11,94 @@
 #  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.
-
+# Source code inspired from https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
 import torchvision
-from timm import create_model
+import torchvision.transforms as transforms
+from torch import nn
 from torch.utils.data import DataLoader
 from torchmetrics.classification import MulticlassAccuracy
-from torchvision import transforms as T
 
 from gradsflow import AutoDataset, Model
-from gradsflow.callbacks import (
-    CometCallback,
-    CSVLogger,
-    EmissionTrackerCallback,
-    ModelCheckpoint,
-    WandbCallback,
-)
-from gradsflow.data.common import random_split_dataset
-
-# Replace dataloaders with your custom dataset and you are all set to train your model
+from gradsflow.callbacks import CSVLogger, ModelCheckpoint
+
+# Replace dataloaders with your custom dataset, and you are all set to train your model
 image_size = (64, 64)
 batch_size = 4
 
-to_rgb = lambda x: x.convert("RGB")
+transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+
+trainset = torchvision.datasets.CIFAR10(root="~/data", train=True, download=True, transform=transform)
+train_dl = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)
 
-augs = T.Compose([to_rgb, T.AutoAugment(), T.Resize(image_size), T.ToTensor()])
-data = torchvision.datasets.CIFAR10("~/data", download=True, transform=augs)
-train_data, val_data = random_split_dataset(data, 0.99)
-train_dl = DataLoader(train_data, batch_size=batch_size)
-val_dl = DataLoader(val_data, batch_size=batch_size)
-num_classes = len(data.classes)
+testset = torchvision.datasets.CIFAR10(root="~/data", train=False, download=True, transform=transform)
+val_dl = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)
+num_classes = len(trainset.classes)
 cbs = [
     CSVLogger(
         verbose=True,
     ),
     ModelCheckpoint(),
-    EmissionTrackerCallback(),
+    # EmissionTrackerCallback(),
     # CometCallback(offline=True),
-    WandbCallback(),
+    # WandbCallback(),
 ]
 
+
+def imshow(img):
+    img = img / 2 + 0.5  # unnormalize
+    npimg = img.numpy()
+    plt.imshow(np.transpose(npimg, (1, 2, 0)))
+    plt.show()
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(3, 6, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        self.fc1 = nn.Linear(16 * 5 * 5, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 10)
+
+    def forward(self, x):
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = torch.flatten(x, 1)  # flatten all dimensions except batch
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+
+
 if __name__ == "__main__":
     autodataset = AutoDataset(train_dl, val_dl, num_classes=num_classes)
-    cnn = create_model("resnet18", pretrained=False, num_classes=num_classes)
+    net = Net()
+    model = Model(net)
+    criterion = nn.CrossEntropyLoss()
+
+    model.compile(
+        criterion,
+        optim.SGD,
+        optimizer_config={"momentum": 0.9},
+        learning_rate=0.001,
+        metrics=[MulticlassAccuracy(autodataset.num_classes)],
+    )
+    model.fit(autodataset, max_epochs=2, callbacks=cbs)
+
+    dataiter = iter(val_dl)
+    images, labels = next(dataiter)
+
+    # print images
+    # imshow(torchvision.utils.make_grid(images))
+    print("GroundTruth: ", " ".join(f"{trainset.classes[labels[j]]:5s}" for j in range(4)))
 
-    model = Model(cnn)
+    outputs = net(images)
+    _, predicted = torch.max(outputs, 1)
 
-    model.compile("crossentropyloss", "adam", metrics=[MulticlassAccuracy(autodataset.num_classes)])
-    model.fit(autodataset, max_epochs=10, steps_per_epoch=10, callbacks=cbs)
+    print("Predicted: ", " ".join(f"{trainset.classes[predicted[j]]:5s}" for j in range(4)))

gradsflow/data/autodata.py

@@ -151,7 +151,7 @@ def _fetch(self, data, device_mapper: Optional[Callable] = None):
         if self.device_setup_status:
             return data
         if device_mapper:
-            data = map(device_mapper, data, self._default_device)
+            data = map(device_mapper, data)
         return data
 
     @property

gradsflow/data/mixins.py

@@ -15,10 +15,13 @@
 
 import torch
 
+from gradsflow.utility import default_device
+
 
 class DataMixin:
     INPUT_KEY = 0  # other common value - inputs, images, text
     OUTPUT_KEY = 1  # other common values - target, ground
+    device = default_device()
 
     def fetch_inputs(self, data: Union[List, Dict]):
         return data[self.INPUT_KEY]

gradsflow/models/base.py

@@ -11,11 +11,11 @@
 #  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 lightning as L
 import smart_open
 import torch
 from lightning.fabric import Fabric
@@ -111,30 +111,39 @@ def __init__(
         self,
         learner: Union[nn.Module, Any],
         device: Optional[str] = "auto",
+        strategy: Optional[str] = None,
+        precision: Any = 32,
+        num_nodes: int = 1,
         use_accelerator: bool = True,
         accelerator_config: dict = None,
     ):
-        self.accelerator = None
+        self._accelerator: L.Fabric = None
         super().__init__()
-        self._set_accelerator(device, use_accelerator, accelerator_config)
+        self._set_accelerator(device, strategy, precision, num_nodes, use_accelerator, accelerator_config)
         self._learner = learner
 
-    def _set_accelerator(self, device: Optional[str], use_accelerate: bool, accelerator_config: dict):
+    def _set_accelerator(
+        self, device: Optional[str], strategy, precision, num_nodes, use_accelerate: bool, accelerator_config: dict
+    ):
         if use_accelerate:
-            self.accelerator = Fabric(accelerator=device, **accelerator_config)
-            self.device = self.accelerator.device
+            self._accelerator = Fabric(
+                accelerator=device, strategy=strategy, precision=precision, num_nodes=num_nodes, **accelerator_config
+            )
+            self.device = self._accelerator.device
         else:
             self.device = device or default_device()
 
     def setup(self, learner: Union[nn.Module, List[nn.Module]], *optimizers):
-        return self.accelerator.setup(learner, *optimizers)
+        if not self._accelerator:
+            return learner, *optimizers
+        return self._accelerator.setup(learner, *optimizers)
 
     def backward(self, loss: torch.Tensor):
         """model.backward(loss)"""
-        if not self.accelerator:
+        if not self._accelerator:
             loss.backward()
         else:
-            self.accelerator.backward(loss)
+            self._accelerator.backward(loss)
 
     def eval(self):
         """Set learner to eval mode for validation"""

gradsflow/models/model.py

@@ -46,6 +46,9 @@ class Model(BaseModel, DataMixin):
 
     Args:
         learner: Trainable model
+        device: auto | cpu | gpu | mps
+        precision: Numerical precision value, could be 32 | 16 | "b16"
+        strategy: Strategy for distributed training (ddp | ddp_spawn | deepspeed | fsdp)
         accelerator_config: Accelerator config
     """
 
@@ -56,13 +59,19 @@ def __init__(
         self,
         learner: Union[nn.Module, Any],
         device: Optional[str] = None,
+        strategy: Optional[str] = None,
+        precision: Any = 32,
+        num_nodes: int = 1,
         use_accelerator: bool = True,
         accelerator_config: dict = None,
     ):
         accelerator_config = accelerator_config or {}
         super().__init__(
             learner=learner,
             device=device,
+            strategy=strategy,
+            precision=precision,
+            num_nodes=num_nodes,
             use_accelerator=use_accelerator,
             accelerator_config=accelerator_config,
         )
@@ -121,6 +130,7 @@ def compile(
             optimizer = optimizer_fn(self.learner.parameters(), lr=learning_rate, **optimizer_config)
         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
@@ -227,7 +237,7 @@ def fit(
         ```python
         autodataset = AutoDataset(train_dataloader, val_dataloader)
         model = Model(cnn)
-        model.compile("crossentropyloss", "adam", learning_rate=1e-3, metrics="accuracy")
+        model.compile("crossentropyloss", "adam", learning_rate=1e-3)
         model.fit(autodataset)
         ```
         Args:
@@ -244,7 +254,7 @@ def fit(
         """
         self.assert_compiled()
         self.autodataset = autodataset
-        self.autodataset.setup_data(self.accelerator)
+        self.autodataset.setup_data(self._accelerator)
 
         if not resume:
             self.tracker.reset()

gradsflow/tuner/automodel.py

@@ -19,7 +19,7 @@
 
 import ray
 from ray import tune
-from ray.tune.sample import Domain
+from ray.tune.search.sample import Domain
 from torch import nn
 
 from gradsflow.data import AutoDataset

tests/models/test_model.py

@@ -89,9 +89,9 @@ def compute_accuracy(*_, **__):
 
 
 def test_set_accelerator(resnet18):
-    model = Model(resnet18, accelerator_config={"precision": 16})
+    model = Model(resnet18, precision=16)
     model.compile()
-    assert model.accelerator
+    assert model._accelerator
 
 
 def test_save_model(tmp_path, resnet18, cnn_model):