remove wandb

benchmarking/classification/__main__.py

@@ -18,7 +18,6 @@
 
 def train_classic(
     problem,
-    trainer: L.Trainer,
     opt: optim.Optimizer,
     hyperparameters: Dict[str, Any],
 ):
@@ -27,20 +26,12 @@ def train_classic(
     data: L.LightningDataModule = problem["dataset"](batch_size=problem["batch_size"])
     classifier = Classifier(problem["model"](), optimizer=opt, **hyperparameters)
 
-    # wandb.init(
-    #     project="pyrfd",
-    #     config={
-    #         "dataset": problem["dataset"].__name__,
-    #         "model": problem["model"].__name__,
-    #         "batch_size": problem["batch_size"],
-    #         "optimizer": opt.__name__,
-    #         "hyperparameters": hyperparameters,
-    #     },
-    # )
+    trainer = trainer_from_problem(problem, opt_name=opt.__name__)
     trainer.fit(classifier, data)
+    trainer.test(classifier, data)
 
 
-def train_rfd(problem, trainer: L.Trainer, cov_string, covariance_model):
+def train_rfd(problem, cov_string, covariance_model):
     """Train a Classifier with RFD"""
     data: L.LightningDataModule = problem["dataset"](batch_size=problem["batch_size"])
     data.prepare_data()
@@ -57,49 +48,41 @@ def train_rfd(problem, trainer: L.Trainer, cov_string, covariance_model):
         problem["model"](), optimizer=RFD, covariance_model=covariance_model
     )
 
-    # wandb.init(
-    #     project="pyrfd",
-    #     config={
-    #         "dataset": problem["dataset"].__name__,
-    #         "model": problem["model"].__name__,
-    #         "batch_size": problem["batch_size"],
-    #         "optimizer": "RFD",
-    #         "covariance": cov_string,
-    #     },
-    # )
+    trainer = trainer_from_problem(problem, opt_name=f"RFD({cov_string})")
     trainer.fit(classifier, data)
+    trainer.test(classifier, data)
 
 
-def main():
-    problem = {
-        "dataset": MNIST,
-        "model": CNN3,
-        "batch_size": 100,
-    }
-
+def trainer_from_problem(problem, opt_name):
     problem_id = f"{problem['dataset'].__name__}_{problem['model'].__name__}_{problem['batch_size']}"
 
-    tlogger = TensorBoardLogger("logs", name=problem_id)
-    csvlogger = CSVLogger("logs", name=problem_id)
+    tlogger = TensorBoardLogger("logs/TensorBoard", name=problem_id + opt_name)
+    csvlogger = CSVLogger("logs", name=problem_id + opt_name)
 
     tlogger.log_hyperparams(params={"batch_size": problem["batch_size"]})
     csvlogger.log_hyperparams(params={"batch_size": problem["batch_size"]})
-    trainer = L.Trainer(
+    return L.Trainer(
         max_epochs=2,
         log_every_n_steps=1,
         logger=[tlogger, csvlogger],
     )
 
+
+def main():
+    problem = {
+        "dataset": MNIST,
+        "model": CNN3,
+        "batch_size": 100,
+    }
+
     train_rfd(
         problem,
-        trainer,
         cov_string="SquaredExponential",
         covariance_model=covariance.SquaredExponential(),
     )
 
     train_classic(
         problem,
-        trainer,
         opt=optim.SGD,
         hyperparameters={
             "lr": 1e-3,
@@ -109,7 +92,6 @@ def main():
 
     train_classic(
         problem,
-        trainer,
         opt=optim.Adam,
         hyperparameters={
             "lr": 1e-3,

benchmarking/classification/cifar/data.py

@@ -0,0 +1,103 @@
+""" CIFAR Data Modules """
+
+import lightning as L
+import torch
+from torch.utils.data import random_split, DataLoader
+
+from torchvision.datasets import CIFAR10 as CIFAR10Dataset, CIFAR100 as CIFAR100Dataset
+from torchvision import transforms
+
+class CIFAR10(L.LightningDataModule):
+    """Represents the CIFAR10 dataset.
+    """
+
+    def __init__(self, data_dir: str=" ./data", batch_size: int = 100) -> None:
+        super().__init__()
+        self.data_dir = data_dir
+        self.batch_size = batch_size
+        # self.transform = transforms.Compose()
+
+    def prepare_data(self) -> None:
+        # download
+        CIFAR10Dataset(self.data_dir, train=True, download=True)
+        CIFAR10Dataset(self.data_dir, train=False, download=True)
+
+    # pylint: disable=attribute-defined-outside-init
+    def setup(self, stage: str):
+        # Assign train/val datasets for use in dataloaders
+        if stage == "fit":
+            cifar10_full = CIFAR10Dataset(self.data_dir, train=True)
+            self.train_data, self.validation_data = random_split(
+                cifar10_full, [0.9, 0.1], generator=torch.Generator().manual_seed(42)
+            )
+
+        # Assign test dataset for use in dataloader(s)
+        if stage == "test":
+            self.test_data = CIFAR10Dataset(
+                self.data_dir, train=False
+            )
+
+        if stage == "predict":
+            self.prediction_data = CIFAR10Dataset(
+                self.data_dir, train=False
+            )
+
+    def train_dataloader(self):
+        return DataLoader(self.train_data, batch_size=self.batch_size, shuffle=True)
+
+    def val_dataloader(self):
+        return DataLoader(self.validation_data, batch_size=self.batch_size)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_data, batch_size=self.batch_size)
+
+    def predict_dataloader(self):
+        return DataLoader(self.prediction_data, batch_size=self.batch_size)
+
+
+class CIFAR100(L.LightningDataModule):
+    """Represents the CIFAR10 dataset.
+    """
+
+    def __init__(self, data_dir: str=" ./data", batch_size: int = 100) -> None:
+        super().__init__()
+        self.data_dir = data_dir
+        self.batch_size = batch_size
+        # self.transform = transforms.Compose()
+
+    def prepare_data(self) -> None:
+        # download
+        CIFAR100Dataset(self.data_dir, train=True, download=True)
+        CIFAR100Dataset(self.data_dir, train=False, download=True)
+
+    # pylint: disable=attribute-defined-outside-init
+    def setup(self, stage: str):
+        # Assign train/val datasets for use in dataloaders
+        if stage == "fit":
+            cifar100_full = CIFAR100Dataset(self.data_dir, train=True)
+            self.train_data, self.validation_data = random_split(
+                cifar100_full, [0.9, 0.1], generator=torch.Generator().manual_seed(42)
+            )
+
+        # Assign test dataset for use in dataloader(s)
+        if stage == "test":
+            self.test_data = CIFAR100Dataset(
+                self.data_dir, train=False
+            )
+
+        if stage == "predict":
+            self.prediction_data = CIFAR100Dataset(
+                self.data_dir, train=False
+            )
+
+    def train_dataloader(self):
+        return DataLoader(self.train_data, batch_size=self.batch_size, shuffle=True)
+
+    def val_dataloader(self):
+        return DataLoader(self.validation_data, batch_size=self.batch_size)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_data, batch_size=self.batch_size)
+
+    def predict_dataloader(self):
+        return DataLoader(self.prediction_data, batch_size=self.batch_size)

benchmarking/classification/classifier.py

@@ -19,7 +19,8 @@ def __init__(self, model, optimizer, **hyperameters):
         self.optimizer = optimizer
         self.model = model
         self.hyperparemeters = hyperameters
-        self.save_hyperparameters()
+        self.save_hyperparameters(ignore=["model", "optimizer", "covariance_model"])
+
 
     # pylint: disable=arguments-differ
     def training_step(self, batch, *args, **kwargs):

benchmarking/classification/mnist/data.py

@@ -36,7 +36,7 @@ def setup(self, stage: str):
         if stage == "fit":
             mnist_full = MNISTDataset(self.data_dir, train=True, transform=self.transform)
             self.train_data, self.validation_data = random_split(
-                mnist_full, [55000, 5000], generator=torch.Generator().manual_seed(42)
+                mnist_full, [0.8, 0.2], generator=torch.Generator().manual_seed(42)
             )
 
         # Assign test dataset for use in dataloader(s)
@@ -91,7 +91,7 @@ def setup(self, stage: str):
         if stage == "fit":
             mnist_full = FashionMNISTDataset(self.data_dir, train=True, transform=self.transform)
             self.train_data, self.validation_data = random_split(
-                mnist_full, [55000, 5000], generator=torch.Generator().manual_seed(42)
+                mnist_full, [0.8, 0.2], generator=torch.Generator().manual_seed(42)
             )
 
         # Assign test dataset for use in dataloader(s)