torchtsetlin.train

torchtsetlin.train.trainer

A PyTorch-style training loop for Tsetlin machines.

Trainer

Trainer(model: TsetlinMachineBase, device=None, callbacks: Optional[Sequence[Callback]] = None, metrics: Optional[Dict[str, Callable]] = None, batch_size: int = 32, eval_batch_size: Optional[int] = None, verbose: bool = True)

Train, validate and test Tsetlin machines with the familiar fit / evaluate / predict trio, PyTorch DataLoader support and callbacks.

Parameters:

Name	Type	Description	Default
model	TsetlinMachineBase	any :class:~torchtsetlin.models.TsetlinMachineBase.	required
device	str | device | None	where to run (defaults to the model's device; the model is moved if given).	None
callbacks	Optional[Sequence[Callback]]	list of :class:~torchtsetlin.train.Callback objects.	None
metrics	Optional[Dict[str, Callable]]	extra metrics {name: fn(votes_or_pred, y) -> float} computed on validation / test data in addition to the defaults (accuracy for classification, MAE/RMSE/R² for regression, Hamming/subset accuracy for multi-label).	None
batch_size	int	default batch size when fit receives tensors or a Dataset.	32
eval_batch_size	Optional[int]	batch size for evaluation (defaults to 4 * batch_size).	None
verbose	bool	print one line per epoch.	True

Example

trainer = Trainer(model, device="cuda", callbacks=[EarlyStopping(patience=5)]) history = trainer.fit((x_train, y_train), epochs=30, val_data=(x_val, y_val)) trainer.evaluate((x_test, y_test))

Source code in src/torchtsetlin/train/trainer.py

def __init__(
    self,
    model: TsetlinMachineBase,
    device=None,
    callbacks: Optional[Sequence[Callback]] = None,
    metrics: Optional[Dict[str, Callable]] = None,
    batch_size: int = 32,
    eval_batch_size: Optional[int] = None,
    verbose: bool = True,
) -> None:
    self.model = model
    if device is not None:
        self.model.to(torch.device(device))
    self.callbacks: List[Callback] = list(callbacks or [])
    if verbose and not any(isinstance(c, ProgressLogger) for c in self.callbacks):
        self.callbacks.append(ProgressLogger())
    for c in self.callbacks:
        c.set_trainer(self)
    self.extra_metrics = dict(metrics or {})
    self.batch_size = int(batch_size)
    self.eval_batch_size = int(eval_batch_size or 4 * batch_size)
    self.history = History()
    self.stop_training = False
    self.epoch = 0
    self.global_step = 0

fit

fit(train_data: DataLike, epochs: int = 10, val_data: Optional[DataLike] = None, batch_size: Optional[int] = None, shuffle: bool = True, eval_train: bool = False) -> History

Train for epochs epochs.

Parameters:

Name	Type	Description	Default
train_data	DataLike	(x, y) tensors/arrays, a Dataset or a DataLoader.	required
epochs	int	number of passes over the data.	10
val_data	Optional[DataLike]	optional validation data evaluated after each epoch (val_* logs).	None
batch_size	Optional[int]	overrides the trainer default for tensor / Dataset inputs.	None
shuffle	bool	shuffle tensor / Dataset inputs each epoch.	True
eval_train	bool	also evaluate on the training data each epoch (train_* logs, slower). The running batch accuracy (batch_accuracy) is always logged.	False

Source code in src/torchtsetlin/train/trainer.py

def fit(
    self,
    train_data: DataLike,
    epochs: int = 10,
    val_data: Optional[DataLike] = None,
    batch_size: Optional[int] = None,
    shuffle: bool = True,
    eval_train: bool = False,
) -> History:
    """Train for ``epochs`` epochs.

    Args:
        train_data: ``(x, y)`` tensors/arrays, a ``Dataset`` or a ``DataLoader``.
        epochs: number of passes over the data.
        val_data: optional validation data evaluated after each epoch (``val_*`` logs).
        batch_size: overrides the trainer default for tensor / ``Dataset`` inputs.
        shuffle: shuffle tensor / ``Dataset`` inputs each epoch.
        eval_train: also evaluate on the training data each epoch (``train_*`` logs,
            slower). The running batch accuracy (``batch_accuracy``) is always logged.
    """
    bs = int(batch_size or self.batch_size)
    batches = self._batches(train_data, bs, shuffle)
    self.stop_training = False
    logs: Dict = {}
    self._call("on_train_begin", logs)
    for epoch in range(epochs):
        self.epoch = epoch
        logs = {}
        self._call("on_epoch_begin", epoch, logs)
        self.model.train()
        if self.model.drop_granularity == "epoch":
            self.model.resample_dropout()
        t0 = time.time()
        correct, total, err_sum = 0.0, 0, 0.0
        for b, (x, y) in enumerate(batches):
            x, y = self._move(x, y)
            self._call("on_batch_begin", b, logs)
            votes = self.model.update(x, y)
            total += votes.shape[0]
            if self.task == "segmentation":
                tgt = self._seg_targets(y, x.shape[0])
                if self._seg_multi_label:
                    flat = tgt.reshape(-1, tgt.shape[-1])
                    correct += float(((votes > 0) == flat.bool()).float().mean(dim=1).sum())
                else:
                    flat = tgt.reshape(-1)
                    ignore = getattr(self.model, "ignore_index", None)
                    keep = flat != ignore if ignore is not None else None
                    pred = votes.argmax(dim=1)
                    if keep is not None:
                        pred, flat = pred[keep], flat[keep]
                        total -= int((~keep).sum())
                    correct += float((pred == flat).sum())
            elif self.task == "classification":
                correct += float((votes.argmax(dim=1) == y.view(-1)).sum())
            elif self.task == "multilabel":
                correct += float(((votes > 0) == y.bool()).float().mean(dim=1).sum())
            else:
                err_sum += float((self.model.votes_to_targets(votes) - y.view(-1)).abs().sum())
            self.global_step += 1
            self._call("on_batch_end", b, logs)
        logs["epoch_time"] = time.time() - t0
        if total:
            if self.task == "regression":
                logs["batch_mae"] = err_sum / total
            else:
                logs["batch_accuracy"] = correct / total
        if eval_train:
            logs.update({f"train_{k}": v for k, v in self.evaluate(train_data, prefix="").items()})
        if val_data is not None:
            logs.update(self.evaluate(val_data, prefix="val_"))
        self._call("on_epoch_end", epoch, logs)
        self.history.append(logs)
        if self.stop_training:
            break
    self._call("on_train_end", logs)
    return self.history

predict

predict(data: DataLike, batch_size: Optional[int] = None, return_votes: bool = False) -> Tensor

Predictions (or vote sums with return_votes=True) for a dataset / tensors.

Source code in src/torchtsetlin/train/trainer.py

@torch.no_grad()
def predict(self, data: DataLike, batch_size: Optional[int] = None, return_votes: bool = False) -> Tensor:
    """Predictions (or vote sums with ``return_votes=True``) for a dataset / tensors."""
    self.model.eval()
    bs = int(batch_size or self.eval_batch_size)
    if self.task == "segmentation" and not return_votes:
        outs = []
        for batch in self._predict_batches(data, bs):
            outs.append(self.model.predict(batch))
        self.model.train()
        return torch.cat(outs, dim=0)
    if isinstance(data, (Tensor, GraphBatch)) or (isinstance(data, (tuple, list)) and len(data) == 2 and not isinstance(data[0], (Dataset, DataLoader))):
        x = data[0] if isinstance(data, (tuple, list)) else data
        x = _as_input(x, self.device)
        outs = [self.model(x[i : i + bs]) for i in range(0, x.shape[0], bs)]
    else:
        outs = []
        for batch in self._batches(data, bs, shuffle=False):
            x = batch[0] if isinstance(batch, (tuple, list)) else batch
            x, _ = self._move(x, None)
            outs.append(self.model(x))
    votes = torch.cat(outs, dim=0)
    if return_votes:
        return votes
    return self._votes_to_pred(votes)

evaluate

evaluate(data: DataLike, batch_size: Optional[int] = None, prefix: str = '') -> Dict[str, float]

Compute metrics over data (model in eval() mode).

Source code in src/torchtsetlin/train/trainer.py

@torch.no_grad()
def evaluate(self, data: DataLike, batch_size: Optional[int] = None, prefix: str = "") -> Dict[str, float]:
    """Compute metrics over ``data`` (model in ``eval()`` mode)."""
    self.model.eval()
    bs = int(batch_size or self.eval_batch_size)
    if self.task == "segmentation":
        out = self._evaluate_segmentation(data, bs)
        self.model.train()
        return {prefix + k: v for k, v in out.items()}
    votes_l, y_l = [], []
    for x, y in self._batches(data, bs, shuffle=False):
        x, y = self._move(x, y)
        votes_l.append(self.model(x))
        y_l.append(y)
    votes = torch.cat(votes_l, dim=0)
    y = torch.cat(y_l, dim=0)
    out: Dict[str, float] = {}
    if self.task == "classification":
        out["accuracy"] = M.accuracy(votes, y.view(-1))
        arg = votes
    elif self.task == "multilabel":
        pred = votes > 0
        out.update(M.multilabel_metrics(pred, y))
        arg = pred
    else:
        pred = self.model.votes_to_targets(votes)
        out.update(M.regression_metrics(pred, y.view(-1)))
        arg = pred
    for name, fn in self.extra_metrics.items():
        out[name] = float(fn(arg, y))
    self.model.train()
    return {prefix + k: v for k, v in out.items()}

test

test(data: DataLike, batch_size: Optional[int] = None) -> Dict[str, float]

Alias of :meth:evaluate with a test_ prefix.

Source code in src/torchtsetlin/train/trainer.py

def test(self, data: DataLike, batch_size: Optional[int] = None) -> Dict[str, float]:
    """Alias of :meth:`evaluate` with a ``test_`` prefix."""
    return self.evaluate(data, batch_size, prefix="test_")

evaluate

evaluate(model: TsetlinMachineBase, data: DataLike, batch_size: int = 512, **kwargs) -> Dict[str, float]

Functional shortcut: Trainer(model, verbose=False).evaluate(data).

Source code in src/torchtsetlin/train/trainer.py

def evaluate(model: TsetlinMachineBase, data: DataLike, batch_size: int = 512, **kwargs) -> Dict[str, float]:
    """Functional shortcut: ``Trainer(model, verbose=False).evaluate(data)``."""
    return Trainer(model, verbose=False, eval_batch_size=batch_size, **kwargs).evaluate(data)

predict

predict(model: TsetlinMachineBase, data, batch_size: int = 512) -> Tensor

Functional shortcut for :meth:Trainer.predict.

Source code in src/torchtsetlin/train/trainer.py

def predict(model: TsetlinMachineBase, data, batch_size: int = 512) -> Tensor:
    """Functional shortcut for :meth:`Trainer.predict`."""
    return Trainer(model, verbose=False, eval_batch_size=batch_size).predict(data)

torchtsetlin.train.callbacks

Callbacks for :class:torchtsetlin.train.Trainer (Keras / Lightning flavoured).

Callback

Base callback. Override any of the hooks; trainer and model are attached by the trainer before training starts.

EarlyStopping

EarlyStopping(monitor: str = 'val_accuracy', patience: int = 10, mode: str = 'max', min_delta: float = 0.0, restore_best: bool = True)

Bases: Callback

Stop training when a monitored metric stops improving.

Parameters:

Name	Type	Description	Default
monitor	str	key in the epoch logs (e.g. "val_accuracy", "val_mae").	'val_accuracy'
patience	int	epochs without improvement before stopping.	10
mode	str	"max" or "min".	'max'
min_delta	float	minimum change to count as an improvement.	0.0
restore_best	bool	reload the best state_dict when training stops.	True

Source code in src/torchtsetlin/train/callbacks.py

def __init__(self, monitor: str = "val_accuracy", patience: int = 10, mode: str = "max", min_delta: float = 0.0, restore_best: bool = True):
    self.monitor = monitor
    self.patience = patience
    self.mode = mode
    self.min_delta = min_delta
    self.restore_best = restore_best
    self.best: float = -math.inf if mode == "max" else math.inf
    self.best_epoch = -1
    self.wait = 0
    self._best_state = None

ModelCheckpoint

ModelCheckpoint(path: str, monitor: str = 'val_accuracy', mode: str = 'max', save_best_only: bool = True, verbose: bool = False)

Bases: Callback

Save model.state_dict() to path whenever the monitored metric improves (save_best_only=True) or after every epoch.

Source code in src/torchtsetlin/train/callbacks.py

def __init__(self, path: str, monitor: str = "val_accuracy", mode: str = "max", save_best_only: bool = True, verbose: bool = False):
    self.path = path
    self.monitor = monitor
    self.mode = mode
    self.save_best_only = save_best_only
    self.verbose = verbose
    self.best: float = -math.inf if mode == "max" else math.inf

CSVLogger

CSVLogger(path: str, append: bool = False)

Bases: Callback

Append the epoch logs to a CSV file.

Source code in src/torchtsetlin/train/callbacks.py

def __init__(self, path: str, append: bool = False):
    self.path = path
    self.append = append
    self._keys = None

LambdaCallback

LambdaCallback(**hooks: Callable)

Bases: Callback

Build a callback from plain functions, e.g. LambdaCallback(on_epoch_end=lambda e, logs: ...).

Source code in src/torchtsetlin/train/callbacks.py

def __init__(self, **hooks: Callable):
    for name, fn in hooks.items():
        if not hasattr(Callback, name):
            raise ValueError(f"Unknown hook {name}")
        setattr(self, name, fn)

HyperparameterSchedule

HyperparameterSchedule(name: str, schedule: Callable[[int], float])

Bases: Callback

Change a model hyper-parameter ("s", "T", "drop_clause_p", ...) at the start of each epoch: schedule(epoch) -> value. Handy for annealing the specificity.

Source code in src/torchtsetlin/train/callbacks.py

def __init__(self, name: str, schedule: Callable[[int], float]):
    self.name = name
    self.schedule = schedule

ProgressLogger

ProgressLogger(print_fn: Callable[[str], None] = print)

Bases: Callback

Print one line per epoch with the logged metrics.

Source code in src/torchtsetlin/train/callbacks.py

def __init__(self, print_fn: Callable[[str], None] = print):
    self.print_fn = print_fn
    self._t0 = 0.0

StateSummaryLogger

Bases: Callback

Add automata statistics (mean/max included literals, empty clauses) to the epoch logs.

torchtsetlin.train.history

Training history container.

History

History()

Per-epoch metrics collected by :class:torchtsetlin.train.Trainer.

history.epochs is a list of dicts; history["val_accuracy"] returns the series of a metric (None where missing); history.best("val_accuracy") the best value/epoch.

Source code in src/torchtsetlin/train/history.py

def __init__(self) -> None:
    self.epochs: List[Dict] = []

to_dataframe

to_dataframe()

Return a pandas.DataFrame (pandas required).

Source code in src/torchtsetlin/train/history.py

def to_dataframe(self):
    """Return a ``pandas.DataFrame`` (pandas required)."""
    import pandas as pd

    return pd.DataFrame(self.epochs)

plot

plot(metrics=None, ax=None)

Plot metric curves (delegates to :func:torchtsetlin.viz.plot_history).

Source code in src/torchtsetlin/train/history.py

def plot(self, metrics=None, ax=None):
    """Plot metric curves (delegates to :func:`torchtsetlin.viz.plot_history`)."""
    from ..viz import plot_history

    return plot_history(self, metrics=metrics, ax=ax)