Merge pull request #3 from PriorLabs/callbacks

Add Callbacks

Author: AlexanderPfefferle

README.md

@@ -62,6 +62,7 @@ from nanotabpfn.train import train
 from nanotabpfn.utils import get_default_device
 from nanotabpfn.interface import NanoTabPFNClassifier
 from torch.nn import CrossEntropyLoss
+from nanotabpfn.callbacks import ConsoleLoggerCallback
 ```
 then we instantiate our model and loss criterion:
 ```python
@@ -81,17 +82,12 @@ prior = PriorDumpDataLoader(filename='50x3_3_100k_classification.h5', num_steps=
 ```
 and finally train our model:
 ```python
-def epoch_callback(epoch, epoch_time, mean_loss, model):
-    classifier = NanoTabPFNClassifier(model, device)
-    # you can add your own eval code here that runs after every epoch
-    print(f'epoch {epoch:5d} | time {epoch_time:5.2f}s | mean loss {mean_loss:5.2f}', flush=True)
-
 trained_model, loss = train(
     model=model,
     prior=prior,
     criterion=criterion,
     epochs=80,
     device=device,
-    epoch_callback=epoch_callback
+    callbacks=[ConsoleLoggerCallback()]
 )
 ```

nanotabpfn/callbacks.py

@@ -0,0 +1,90 @@
+from abc import ABC, abstractmethod
+
+
+class Callback(ABC):
+    """ Abstract base class for callbacks."""
+
+    @abstractmethod
+    def on_epoch_end(self, epoch: int, epoch_time: float, loss: float, model, **kwargs):
+        """
+        Called at the end of each epoch.
+
+        Args:
+            epoch (int): The current epoch number.
+            epoch_time (float): Time of the epoch in seconds.
+            loss (float): Mean loss for the epoch.
+            model: The model being trained.
+            **kwargs: Additional arguments.
+        """
+        pass
+
+    @abstractmethod
+    def close(self):
+        """
+        Called to release any resources or perform cleanup.
+        """
+        pass
+
+
+class BaseLoggerCallback(Callback):
+    """ Abstract base class for logger callbacks. """
+    pass
+
+
+class ConsoleLoggerCallback(BaseLoggerCallback):
+    """ Logger callback that prints epoch information to the console. """
+
+    def on_epoch_end(self, epoch: int, epoch_time: float, loss: float, model, **kwargs):
+        print(f'Epoch {epoch:5d} | Time {epoch_time:5.2f}s | Mean Loss {loss:5.2f}', flush=True)
+
+    def close(self):
+        """ Nothing to clean up for print logger. """
+        pass
+
+
+class TensorboardLoggerCallback(BaseLoggerCallback):
+    """ Logger callback that logs epoch information to TensorBoard. """
+
+    def __init__(self, log_dir: str):
+        from torch.utils.tensorboard import SummaryWriter
+        self.writer = SummaryWriter(log_dir=log_dir)
+
+    def on_epoch_end(self, epoch: int, epoch_time: float, loss: float, model, **kwargs):
+        self.writer.add_scalar('Loss/train', loss, epoch)
+        self.writer.add_scalar('Time/epoch', epoch_time, epoch)
+
+    def close(self):
+        self.writer.close()
+
+
+class WandbLoggerCallback(BaseLoggerCallback):
+    """ Logger callback that logs epoch information to Weights & Biases. """
+
+    def __init__(self, project: str, name: str = None, config: dict = None, log_dir: str = None):
+        """
+        Initializes a WandbLoggerCallback.
+
+        Args:
+            project (str): The name of the wandb project.
+            name (str, optional): The name of the run. Defaults to None.
+            config (dict, optional): Configuration dictionary for the run. Defaults to None.
+            log_dir (str, optional): Directory to save wandb logs. Defaults to None.
+        """
+        try:
+            import wandb
+            self.wandb = wandb  # store wandb module to avoid import if not used
+            wandb.init(
+                project=project,
+                name=name,
+                config=config,
+                dir=log_dir,
+            )
+        except ImportError:
+            raise ImportError("wandb is not installed. Install it with: pip install wandb") from e
+
+    def on_epoch_end(self, epoch: int, epoch_time: float, loss: float, model, **kwargs):
+        log_dict = {'epoch': epoch, 'loss': loss, ' epoch_time': epoch_time}
+        self.wandb.log(log_dict)
+
+    def close(self):
+        self.wandb.finish()

nanotabpfn/train.py

@@ -2,16 +2,18 @@
 from torch import nn
 import time
 from torch.utils.data import DataLoader
-from typing import Tuple, Dict, Callable
+from typing import Dict
 from pfns.bar_distribution import FullSupportBarDistribution
 import schedulefree
 
+from nanotabpfn.callbacks import Callback
 from nanotabpfn.model import NanoTabPFNModel
 from nanotabpfn.utils import get_default_device
 
-def train(model: NanoTabPFNModel, prior: DataLoader, criterion: nn.CrossEntropyLoss | FullSupportBarDistribution, epochs: int,
-          accumulate_gradients: int = 1, lr: float = 1e-4, device: torch.device = None,
-	  epoch_callback: Callable[[int, float, float, NanoTabPFNModel, FullSupportBarDistribution | None], None] = None, ckpt: Dict[str, torch.Tensor] = None):
+
+def train(model: NanoTabPFNModel, prior: DataLoader, criterion: nn.CrossEntropyLoss | FullSupportBarDistribution,
+          epochs: int, accumulate_gradients: int = 1, lr: float = 1e-4, device: torch.device = None,
+          callbacks: list[Callback]=None, ckpt: Dict[str, torch.Tensor] = None):
     """
     Trains our model on the given prior using the given criterion.
 
@@ -22,14 +24,17 @@ def train(model: NanoTabPFNModel, prior: DataLoader, criterion: nn.CrossEntropyL
         epochs: (int) the number of epochs we train for, the number of steps that constitute an epoch are decided by the prior
         accumulate_gradients: (int) the number of gradients to accumulate before updating the weights
         device: (torch.device) the device we are using
-        epoch_callback: (Callable[[int, float, float, NanoTabPFNModel], None]) optional callback function that will be called
-	                at the end of each epoch with the current epoch, epoch duration, mean loss, and the model,
-			intended to be used for logging/validation/evaluation
+        callbacks: A list of callback instances to execute at the end of each epoch. These can be used for
+            logging, validation, or other custom actions.
+        ckpt (Dict[str, torch.Tensor], optional): A checkpoint dictionary containing the model and optimizer states,
+            as well as the last completed epoch. If provided, training resumes from this checkpoint.
 
     Returns:
         (torch.Tensor) a tensor of shape (num_rows, batch_size, num_features, embedding_size)
     """
     # print(f"Using a Transformer with {sum(p.numel() for p in model.parameters())/1000/1000:.{2}f} M parameters")
+    if callbacks is None:
+        callbacks = []
     if not device:
         device = get_default_device()
     model.to(device)
@@ -41,8 +46,8 @@ def train(model: NanoTabPFNModel, prior: DataLoader, criterion: nn.CrossEntropyL
     assert prior.num_steps % accumulate_gradients == 0, 'num_steps must be divisible by accumulate_gradients'
 
     try:
-        for epoch in range(ckpt['epoch']+1 if ckpt else 1, epochs + 1):
-            start_time = time.time()
+        for epoch in range(ckpt['epoch'] + 1 if ckpt else 1, epochs + 1):
+            epoch_start_time = time.time()
             model.train()  # Turn on the train mode
             optimizer.train()
             total_loss = 0.
@@ -81,12 +86,15 @@ def train(model: NanoTabPFNModel, prior: DataLoader, criterion: nn.CrossEntropyL
             }
             torch.save(training_state, 'latest_checkpoint.pth')
 
-            if epoch_callback:
+            for callback in callbacks:
                 if type(criterion) is FullSupportBarDistribution:
-                    epoch_callback(epoch, end_time - start_time, mean_loss, model, dist=criterion)
+                    callback.on_epoch_end(epoch, end_time - epoch_start_time, mean_loss, model, dist=criterion)
                 else:
-                    epoch_callback(epoch, end_time-start_time, mean_loss, model)
+                    callback.on_epoch_end(epoch, end_time - epoch_start_time, mean_loss, model)
     except KeyboardInterrupt:
         pass
+    finally:
+        for callback in callbacks:
+            callback.close()
 
     return model, total_loss

pretrain_classification.py

@@ -5,6 +5,7 @@
 from torch import nn
 from functools import partial
 
+from nanotabpfn.callbacks import ConsoleLoggerCallback
 from nanotabpfn.priors import PriorDumpDataLoader
 from nanotabpfn.model import NanoTabPFNModel
 from nanotabpfn.train import train
@@ -29,15 +30,14 @@
 parser.add_argument("-epochs", type=int, default=10000, help="number of epochs to train for")
 parser.add_argument("-loadcheckpoint", type=str, default=None, help="checkpoint from which to continue training")
 
-
 args = parser.parse_args()
 
 set_randomness_seed(2402)
 
 device = get_default_device()
-ckpt=None
+ckpt = None
 if args.loadcheckpoint:
-    ckpt=torch.load(args.loadcheckpoint)
+    ckpt = torch.load(args.loadcheckpoint)
 
 prior = PriorDumpDataLoader(filename=args.priordump, num_steps=args.steps, batch_size=args.batchsize, device=device, starting_index=args.steps*(ckpt['epoch'] if ckpt else 0))
 
@@ -60,17 +60,24 @@
 datasets.append(train_test_split(*load_wine(return_X_y=True), test_size=0.5, random_state=42))
 datasets.append(train_test_split(*load_breast_cancer(return_X_y=True), test_size=0.5, random_state=42))
 
-def epoch_callback(epoch, epoch_time, mean_loss, model):
-    classifier = NanoTabPFNClassifier(model, device)
-    scores = []
-    for  X_train, X_test, y_train, y_test in datasets:
-        classifier.fit(X_train, y_train)
-        pred = classifier.predict(X_test)
-        scores.append(accuracy_score(y_test, pred))
-    avg_score = sum(scores)/len(scores)
-    print(f'epoch {epoch:5d} | time {epoch_time:5.2f}s | mean loss {mean_loss:5.2f} | avg accuracy {avg_score:.3f}', flush=True)
+
+class EvaluationLoggerCallback(ConsoleLoggerCallback):
+    def __init__(self, datasets):
+        self.datasets = datasets
+
+    def on_epoch_end(self, epoch: int, epoch_time: float, loss: float, model, **kwargs):
+        classifier = NanoTabPFNClassifier(model, device)
+        scores = []
+        for X_train, X_test, y_train, y_test in self.datasets:
+            classifier.fit(X_train, y_train)
+            pred = classifier.predict(X_test)
+            scores.append(accuracy_score(y_test, pred))
+        avg_score = sum(scores) / len(scores)
+        print(f'epoch {epoch:5d} | time {epoch_time:5.2f}s | mean loss {loss:5.2f} | avg accuracy {avg_score:.3f}',
+              flush=True)
 
 
+callbacks = [EvaluationLoggerCallback(datasets)]
 
 trained_model, loss = train(
     model=model,
@@ -80,7 +87,7 @@ def epoch_callback(epoch, epoch_time, mean_loss, model):
     accumulate_gradients=args.accumulate,
     lr=args.lr,
     device=device,
-    epoch_callback=epoch_callback,
+    callbacks=callbacks,
     ckpt=ckpt
 )
 

pretrain_regression.py

@@ -1,6 +1,7 @@
 import argparse
 import torch
 
+from nanotabpfn.callbacks import ConsoleLoggerCallback
 from nanotabpfn.priors import PriorDumpDataLoader
 from nanotabpfn.model import NanoTabPFNModel
 from nanotabpfn.train import train
@@ -30,15 +31,14 @@
 parser.add_argument("-loadcheckpoint", type=str, default=None, help="checkpoint from which to continue training")
 parser.add_argument("-n_buckets", type=int, default=100, help="number of buckets for the data loader")
 
-
 args = parser.parse_args()
 
 set_randomness_seed(2402)
 
 device = get_default_device()
-ckpt=None
+ckpt = None
 if args.loadcheckpoint:
-    ckpt=torch.load(args.loadcheckpoint)
+    ckpt = torch.load(args.loadcheckpoint)
 
 prior = PriorDumpDataLoader(filename=args.priordump, num_steps=args.steps, batch_size=args.batchsize, device=device, starting_index=args.steps*(ckpt['epoch'] if ckpt else 0))
 
@@ -69,17 +69,24 @@
 datasets = []
 datasets.append(train_test_split(*load_diabetes(return_X_y=True), test_size=0.5, random_state=42))
 
-def epoch_callback(epoch, epoch_time, mean_loss, model, dist):
-    regressor = NanoTabPFNRegressor(model, dist, device)
-    scores = []
-    for  X_train, X_test, y_train, y_test in datasets:
-        regressor.fit(X_train, y_train)
-        pred = regressor.predict(X_test)
-        scores.append(r2_score(y_test, pred))
-        print(r2_score(y_test, pred))
-    avg_score = sum(scores)/len(scores)
-    print(f'epoch {epoch:5d} | time {epoch_time:5.2f}s | mean loss {mean_loss:5.2f} | avg r2 score {avg_score:.3f}', flush=True)
 
+class EvaluationLoggerCallback(ConsoleLoggerCallback):
+    def __init__(self, datasets):
+        self.datasets = datasets
+
+    def on_epoch_end(self, epoch: int, epoch_time: float, loss: float, model, **kwargs):
+        regressor = NanoTabPFNRegressor(model, dist, device)
+        scores = []
+        for X_train, X_test, y_train, y_test in datasets:
+            regressor.fit(X_train, y_train)
+            pred = regressor.predict(X_test)
+            scores.append(r2_score(y_test, pred))
+        avg_score = sum(scores) / len(scores)
+        print(f'epoch {epoch:5d} | time {epoch_time:5.2f}s | mean loss {loss:5.2f} | avg r2 score {avg_score:.3f}',
+              flush=True)
+
+
+callbacks = [EvaluationLoggerCallback(datasets)]
 
 trained_model, loss = train(
     model=model,
@@ -89,7 +96,7 @@ def epoch_callback(epoch, epoch_time, mean_loss, model, dist):
     accumulate_gradients=args.accumulate,
     lr=args.lr,
     device=device,
-    epoch_callback=epoch_callback,
+    callbacks=callbacks,
     ckpt=ckpt
 )
 

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 name = "nanotabpfn"
 version = "0.0.1"
 authors = [
-  { name="Alexander Pfefferle", email="[email protected]" },
+    { name = "Alexander Pfefferle", email = "[email protected]" },
 ]
 description = "A Playground for Tabular Foundation Models"
 readme = "README.md"
@@ -25,6 +25,9 @@ dependencies = [
     "pfns==0.3.0",
 ]
 
+[project.optional-dependencies]
+wandb = ["wandb>=0.20"]
+tensorboard = ["tensorboard>=2.19"]
 [project.urls]
 Homepage = "https://github.com/PriorLabs/nanoTabPFN"
 Issues = "https://github.com/PriorLabs/nanoTabPFN/issues"