Classification Trainer

class hivegraph.engine.classification.ClassificationTrainer(model: ~torch.nn.modules.module.Module, dataset: ~torch_geometric.data.dataset.Dataset, criterion: ~typing.Callable = <function nll_loss>, device: str = 'cpu', num_folds: int = 10, random_state: int = 42, test_metric: str | None = None)

Bases: BaseTrainer

Custom Trainer Class for Classification Tasks

fit(epochs: int, batch_size: int, optimizer: Optimizer, verbose: bool = True, log_to_wandb: bool = False) → None

Trains the model on the given dataset

Parameters:

• epochs (int) – Number of epochs

• batch_size (int) – batch size

• optimizer (torch.optim.Optimizer) – Optimizer

• verbose (bool) – Whether to print stats, defaults to True

• log_to_wandb (bool) – Whether to log stats to wandb, defaults to False

train_step(train_dataloder: Iterable, optimizer: Optimizer) → float

Performs a single training step

Parameters:

• train_dataloder (Iterable) – Training Dataloader

• optimizer (torch.optim.Optimizer) – Optimizer

Returns:

Training Loss

Return type:

float

val_step(val_dataloader: Iterable) → float

Calculates the loss on the validation set

Parameters:

val_dataloader (Iterable) – Validation Dataloader

Returns:

Validation Loss

Return type:

float

test_step(test_dataloader: Iterable, test_metric: str | None) → float

Calculates the accuracy on the test set

Parameters:

• test_dataloader (Iterable) – Test Dataloader

• test_metric (str) – Test Metric, either “Accuracy” or “AUROC”

Returns:

Test Accuracy

Return type:

float