Training API

DINOTrainer

DINOTrainer(grad_clip_norm: Optional[float] = 3.0, device: device | str = 'cpu', callbacks: Optional[List[Callback]] = None)

Stateless trainer compatible with any BaseSSLModel.

Per-batch order: compute_loss → backward → clip_grad → post_backward → optimizer.step → post_step Per-epoch order: on_epoch_start → batches → on_epoch_end

Parameters:

Name	Type	Description	Default
grad_clip_norm	Optional[float]	Clip student gradient norm. ~3.0 stabilises early DINO.	3.0
device	device | str	Default device used in train() when none is given per call.	'cpu'
callbacks	Optional[List[Callback]]	List of Callback instances invoked at epoch/batch boundaries.	None

Source code in src/graphssl/training/trainer.py

def __init__(
    self,
    grad_clip_norm: Optional[float] = 3.0,
    device: torch.device | str = "cpu",
    callbacks: Optional[List[Callback]] = None,
):
    self.grad_clip_norm = grad_clip_norm
    self.device = torch.device(device)
    self.callbacks: List[Callback] = callbacks or []

train_epoch

train_epoch(model: BaseSSLModel, loader: DataLoader, optimizer: Optimizer, device: device | str | None = None, epoch: int = 0) -> float

Run one training epoch and return the average loss.

Source code in src/graphssl/training/trainer.py

def train_epoch(
    self,
    model: BaseSSLModel,
    loader: DataLoader,
    optimizer: Optimizer,
    device: torch.device | str | None = None,
    epoch: int = 0,
) -> float:
    """Run one training epoch and return the average loss."""
    _device = torch.device(device) if device is not None else self.device
    model.train()
    total_loss = 0.0
    n_batches = 0

    for batch in loader:
        batch = batch.to(_device)
        loss = model.compute_loss(batch)

        optimizer.zero_grad()
        loss.backward()

        if self.grad_clip_norm is not None:
            torch.nn.utils.clip_grad_norm_(
                list(model.student_parameters()),
                self.grad_clip_norm,
            )

        model.post_backward()
        optimizer.step()
        model.post_step()

        for cb in self.callbacks:
            cb.on_batch_end(self, model, loss, batch)

        total_loss += loss.item()
        n_batches += 1

    return total_loss / max(n_batches, 1)

train

train(model: BaseSSLModel, loader: DataLoader, optimizer: Optimizer, num_epochs: int, device: device | str | None = None) -> List[float]

Train for num_epochs epochs, fire callbacks, return per-epoch losses.

Source code in src/graphssl/training/trainer.py

def train(
    self,
    model: BaseSSLModel,
    loader: DataLoader,
    optimizer: Optimizer,
    num_epochs: int,
    device: torch.device | str | None = None,
) -> List[float]:
    """Train for num_epochs epochs, fire callbacks, return per-epoch losses."""
    _device = torch.device(device) if device is not None else self.device
    model = model.to(_device)

    for cb in self.callbacks:
        cb.on_train_start(self, model)

    losses: List[float] = []
    for epoch in range(num_epochs):
        model.on_epoch_start(epoch)
        for cb in self.callbacks:
            cb.on_epoch_start(self, model, epoch)

        avg_loss = self.train_epoch(model, loader, optimizer, device=_device, epoch=epoch)
        losses.append(avg_loss)
        metrics = {"loss": avg_loss}

        model.on_epoch_end(epoch)
        for cb in self.callbacks:
            cb.on_epoch_end(self, model, epoch, metrics)

    for cb in self.callbacks:
        cb.on_train_end(self, model)

    return losses

DataModule

DataModule(data: Optional[Data] = None, labels: Optional[Tensor] = None, train_idx: Optional[Tensor] = None, val_idx: Optional[Tensor] = None, test_idx: Optional[Tensor] = None, is_graph_level: bool = False, batch_size: int = 32, dataset_obj: Optional[Dataset] = None, num_workers: int = 0)

Wraps a dataset and exposes train/eval DataLoaders.

Node-level task (is_graph_level=False): - data: single PyG Data object with train_mask / val_mask / test_mask encoded as index tensors (train_idx, val_idx, test_idx). - train_dataloader() returns a single-element iterable yielding the raw full-graph Data object (not routed through PyG's Batch collation, so no spurious .batch_size is attached). - neighbor_loader() returns a NeighborLoader for mini-batch training.

Graph-level task (is_graph_level=True): - dataset_obj: a PyG Dataset or list of Data objects. - train_dataloader() / eval_dataloader() iterate over the dataset.

Attributes exposed for evaluation utilities

data, labels, train_idx, val_idx, test_idx, is_graph_level, batch_size, dataset_obj.

Source code in src/graphssl/data/datamodule.py

def __init__(
    self,
    data: Optional[Data] = None,
    labels: Optional[torch.Tensor] = None,
    train_idx: Optional[torch.Tensor] = None,
    val_idx: Optional[torch.Tensor] = None,
    test_idx: Optional[torch.Tensor] = None,
    is_graph_level: bool = False,
    batch_size: int = 32,
    dataset_obj: Optional[Dataset] = None,
    num_workers: int = 0,
):
    self.data = data
    self.labels = labels
    self.train_idx = train_idx
    self.val_idx = val_idx
    self.test_idx = test_idx
    self.is_graph_level = is_graph_level
    self.batch_size = batch_size
    self.dataset_obj = dataset_obj
    self.num_workers = num_workers

neighbor_loader

neighbor_loader(num_neighbors: List[int], input_nodes: Optional[Tensor] = None, **kwargs) -> NeighborLoader

Mini-batch NeighborLoader for node-level training.

Source code in src/graphssl/data/datamodule.py

def neighbor_loader(
    self,
    num_neighbors: List[int],
    input_nodes: Optional[torch.Tensor] = None,
    **kwargs,
) -> NeighborLoader:
    """Mini-batch NeighborLoader for node-level training."""
    assert self.data is not None, "data must be set for node-level mini-batch loading"
    return NeighborLoader(
        self.data,
        num_neighbors=num_neighbors,
        input_nodes=input_nodes,
        batch_size=self.batch_size,
        num_workers=self.num_workers,
        **kwargs,
    )