""" Shared machinery for QuaPy's "bag-trained" neural quantifiers, i.e., methods that -- like :class:`quapy.method.meta.QuaNet` -- do not follow the classify-then-aggregate pattern of :class:`quapy.method.aggregative.AggregativeQuantifier`, but are instead trained and evaluated end-to-end on whole samples ("bags") of known prevalence. :class:`quapy.method._histnet.HistNetQ` and :class:`quapy.method._gmnet.GMNet` share the same overall architecture (feature_extraction -> quantification module -> small MLP head -> softmax/normalize) and the same bag-based training protocol (bag generation via a QuaPy sampling protocol, early stopping, LR scheduling, checkpointing). This module factors that common part out into :class:`BagTrainedQuantifier`; concrete subclasses only need to supply the quantification module placed between the feature extractor and the shared head (see :meth:`BagTrainedQuantifier._build_quantmodule`). """ import copy import os import random from abc import ABC, abstractmethod import numpy as np import torch import torch.nn as nn from tqdm import tqdm from quapy.data import LabelledCollection from quapy.method.base import BaseQuantifier from quapy.protocol import AbstractProtocol, UPP from quapy.util import EarlyStop class IdentityFeatureExtractionModule(nn.Module): """Used when no feature extraction module is provided: instances are assumed to already be in their final numeric representation.""" def __init__(self, input_size): super().__init__() self.output_size = input_size def forward(self, x): return x def to_tensor(x, device): if torch.is_tensor(x): return x.to(device=device, dtype=torch.float32) if hasattr(x, 'toarray'): # scipy sparse x = x.toarray() return torch.as_tensor(np.asarray(x), dtype=torch.float32, device=device) def stack_bags(bags, device): """ :param bags: an iterable of (X_bag, prevalence) pairs, all X_bag with the same number of instances :return: a pair of tensors (X, P) of shape (n_bags, bag_size, n_features) and (n_bags, n_classes) """ Xs, ps = zip(*bags) X = torch.stack([to_tensor(x, device) for x in Xs]) P = torch.stack([to_tensor(p, device) for p in ps]) return X, P def mix_two_bags(bag_a, bag_b, bag_size, rng): """Synthesizes a new bag of size `bag_size` by mixing two given bags with a random ratio, following the "mixer" idea from the original HistNetQ repo (`UnlabeledMixerBagGenerator`): useful when the only available training material is a modest number of pre-built samples (e.g., LeQua's dev samples) and one wants extra intermediate-prevalence bags without access to instance-level labels. """ Xa, pa = bag_a Xb, pb = bag_b m = rng.random() na = round(m * bag_size) nb = bag_size - na idx_a = rng.choices(range(len(Xa)), k=na) if na > 0 else [] idx_b = rng.choices(range(len(Xb)), k=nb) if nb > 0 else [] Xa, Xb = np.asarray(Xa), np.asarray(Xb) X_mixed = np.concatenate([Xa[idx_a], Xb[idx_b]], axis=0) p_mixed = m * np.asarray(pa, dtype=float) + (1 - m) * np.asarray(pb, dtype=float) return X_mixed, p_mixed def build_output_head(input_size, n_classes, linear_sizes, dropout, output_function): """Builds the small MLP + output activation shared by every bag-trained quantifier's head: a stack of (Linear, LeakyReLU, Dropout) blocks sized by `linear_sizes`, followed by a final Linear to `n_classes` and either a softmax or an L1-normalization (applied in :class:`BagNetworkModule`), both yielding a valid prevalence vector.""" output_module = nn.Sequential() prev_size = input_size for i, linear_size in enumerate(linear_sizes): output_module.add_module(f'linear_{i}', nn.Linear(prev_size, linear_size)) output_module.add_module(f'leakyrelu_{i}', nn.LeakyReLU()) output_module.add_module(f'dropout_{i}', nn.Dropout(dropout)) prev_size = linear_size output_module.add_module('last_linear', nn.Linear(prev_size, n_classes)) if output_function == 'softmax': output_module.add_module('softmax', nn.Softmax(dim=1)) elif output_function == 'normalize': output_module.add_module('relu', nn.ReLU()) else: raise ValueError(f"unknown {output_function=}; valid ones are 'softmax', 'normalize'") return output_module class BagNetworkModule(nn.Module): """The full network shared by every bag-trained quantifier: feature extraction, a pluggable quantification module (mapping a bag of instance-level features to a single per-bag representation), and the shared MLP head. :param quantmodule: a `torch.nn.Module` exposing an `output_size` attribute, mapping a tensor of shape (batch_size, bag_size, n_features) to one of shape (batch_size, quantmodule.output_size). """ def __init__(self, feature_extraction_module, quantmodule, n_classes, linear_sizes, dropout, output_function): super().__init__() self.feature_extraction_module = feature_extraction_module self.quantmodule = quantmodule self.output_function = output_function self.output_module = build_output_head( quantmodule.output_size, n_classes, linear_sizes, dropout, output_function ) def forward(self, bag): # bag: (batch_size, bag_size, n_features) features = self.feature_extraction_module(bag) representation = self.quantmodule(features) out = self.output_module(representation) if self.output_function == 'normalize': out = nn.functional.normalize(out, p=1, dim=1) return out class BagTrainedQuantifier(BaseQuantifier, ABC): """ Base class for QuaPy's neural quantifiers trained end-to-end on whole samples ("bags") of known prevalence, rather than following the classify-then-aggregate pattern of :class:`quapy.method.aggregative.AggregativeQuantifier` (in the spirit of :class:`quapy.method.meta.QuaNet`). Concrete subclasses only need to provide the quantification module placed between the feature extractor and the shared MLP head (see :meth:`_build_quantmodule`) and a checkpoint-name prefix (see :attr:`_checkpoint_prefix`); bag generation, the training/validation loop, early stopping, LR scheduling, checkpointing, and prediction are all shared. Training data can be provided in two ways: * via :meth:`fit`, from a plain labelled collection (`X`, `y`): training/validation bags are then generated by resampling from it using a QuaPy sampling protocol (:class:`quapy.protocol.UPP` by default). * via :meth:`fit_from_samples`, from a :class:`quapy.protocol.AbstractProtocol` that already yields the training bags (e.g., :class:`quapy.data._lequa.SamplesFromDir` for LeQua-style pre-built samples), optionally enriched with synthetic bags mixed from the given ones. :param feature_extraction_module: a `torch.nn.Module` exposing an `output_size` attribute, used to embed each instance before the quantification module (e.g., a small MLP for tabular data, a CNN for images). If None (default), an identity module is used, i.e., the instances in `X` are assumed to already be in their final numeric representation. :param linear_sizes: tuple of ints with the sizes of the linear layers used in the shared head (default empty, i.e., only the final classification layer is used). :param dropout: dropout applied after each of the `linear_sizes` layers (default 0). :param output_function: either 'softmax' or 'normalize' (L1); both yield a valid prevalence vector (default 'softmax'). :param bag_size: number of instances per training/validation bag (default 500). :param n_bags_train: number of bags generated per training epoch (default 500). :param n_bags_val: number of bags generated per validation epoch (default 500). :param train_epochs: maximum number of training epochs (default 200). :param patience: number of epochs without improvement in validation loss before early-stopping (default 20). :param start_lr: initial learning rate (default 1e-3). :param end_lr: once the learning rate decays below this value, training stops (default 1e-6). :param lr_factor: factor by which the learning rate is reduced after `patience` epochs without improvement (default 0.1). :param weight_decay: L2 regularization (default 0). :param quant_loss: the quantification loss to minimize (default `torch.nn.L1Loss()`), called as `quant_loss(true_prevalences, predicted_prevalences)`. :param batch_size: number of bags per gradient update (default 16). :param protocol: the :class:`quapy.protocol.AbstractStochasticSeededProtocol` subclass used by :meth:`fit` to resample bags from the given labelled collection (default :class:`quapy.protocol.UPP`, which draws bags with prevalence sampled uniformly at random from the simplex). :param protocol_params: dict of extra keyword arguments passed to `protocol` (besides `data`, `sample_size`, `repeats`, and `random_state`, which are set internally); default None. :param val_split: float in (0,1), the proportion of the collection given to :meth:`fit` that is held out (via stratified sampling) for validation and early stopping (default 0.4). :param device: `'cpu'` or `'cuda'` (default 'cpu'). :param random_state: seed used for the train/validation split and for the (fixed) validation sampling sequence (default 0). :param checkpointdir: directory where the best model found during training is stored (default '../checkpoint'). :param checkpointname: name of the checkpoint file; if None (default), a random name prefixed by :attr:`_checkpoint_prefix` is generated. :param verbose: verbosity level; if >0, shows a progress bar with the current losses (default 0). """ def __init__(self, feature_extraction_module=None, linear_sizes=(), dropout=0., output_function='softmax', bag_size=500, n_bags_train=500, n_bags_val=500, train_epochs=200, patience=20, start_lr=1e-3, end_lr=1e-6, lr_factor=0.1, weight_decay=0., quant_loss=None, batch_size=16, protocol=UPP, protocol_params=None, val_split=0.4, device='cpu', random_state=0, checkpointdir='../checkpoint', checkpointname=None, verbose=0): self.feature_extraction_module = feature_extraction_module self.linear_sizes = linear_sizes self.dropout = dropout self.output_function = output_function self.bag_size = bag_size self.n_bags_train = n_bags_train self.n_bags_val = n_bags_val self.train_epochs = train_epochs self.patience = patience self.start_lr = start_lr self.end_lr = end_lr self.lr_factor = lr_factor self.weight_decay = weight_decay self.quant_loss = quant_loss if quant_loss is not None else torch.nn.L1Loss() self.batch_size = batch_size self.protocol = protocol self.protocol_params = protocol_params self.val_split = val_split self.device = torch.device(device) self.random_state = random_state if checkpointname is None: local_random = random.Random() random_code = '-'.join(str(local_random.randint(0, 1000000)) for _ in range(5)) checkpointname = f'{self._checkpoint_prefix}-{random_code}' self.checkpointdir = checkpointdir self.checkpoint = os.path.join(checkpointdir, checkpointname) self.verbose = verbose self._classes_ = None self.model = None @property def classes_(self): return self._classes_ @property @abstractmethod def _checkpoint_prefix(self): """Short name used as the default checkpoint filename prefix (e.g. 'HistNetQ', 'GMNet').""" ... @abstractmethod def _build_quantmodule(self, n_features): """Builds the module placed between the (already feature-extracted) instances and the shared MLP head. Must expose an `output_size` attribute and accept input of shape (batch_size, bag_size, n_features), returning one of shape (batch_size, output_size).""" ... def _extra_loss(self): """Optional additional term added to the quantification loss during training (e.g., GMNet's CKA regularization across GM layers). Returns 0 by default.""" return 0. def fit(self, X, y): """ Trains the quantifier from a plain labelled collection, generating training and validation bags by resampling from it via `self.protocol` (a fresh random sequence of bags every epoch for training, and a fixed, reproducible sequence for validation). :param X: the training instances :param y: the labels of X :return: self """ data = LabelledCollection(X, y) self._classes_ = data.classes_ train_data, val_data = data.split_stratified(train_prop=1 - self.val_split, random_state=self.random_state) protocol_params = self.protocol_params or {} def train_bags(): sampler = self.protocol( train_data, sample_size=self.bag_size, repeats=self.n_bags_train, random_state=None, **protocol_params ) return sampler() def val_bags(): sampler = self.protocol( val_data, sample_size=self.bag_size, repeats=self.n_bags_val, random_state=self.random_state, **protocol_params ) return sampler() n_features = train_data.instances.shape[1] self._fit_loop(train_bags, val_bags, n_features, n_bags_train=self.n_bags_train, n_bags_val=self.n_bags_val) return self def fit_from_samples(self, protocol: AbstractProtocol, val_protocol: AbstractProtocol = None, mix_bags=False, mix_bags_proportion=0.5): """ Trains the quantifier from a protocol that already yields the training bags (e.g., :class:`quapy.data._lequa.SamplesFromDir`, for LeQua-style pre-built samples), instead of resampling from a labelled collection. This is the entry point to use whenever only bags of known prevalence are available (no instance-level labels). :param protocol: an :class:`AbstractProtocol` yielding `(sample, prevalence)` pairs; consumed once and kept in memory (expected to be of modest size, as is typical of pre-built sample collections). :param val_protocol: an optional, separate protocol providing the validation bags; if None, a `val_split` fraction of the bags returned by `protocol` is held out instead. :param mix_bags: if True, in addition to the bags returned by `protocol`, synthesize extra bags each epoch by mixing random pairs of the given bags with a random ratio (a substitute for the original HistNetQ repo's `UnlabeledMixerBagGenerator`, useful to broaden the coverage of prevalence values beyond what the given bags exhibit). :param mix_bags_proportion: proportion (relative to the number of base training bags) of extra mixed bags to generate per epoch when `mix_bags=True` (default 0.5). :return: self """ assert isinstance(protocol, AbstractProtocol), 'protocol must be an instance of AbstractProtocol' base_bags = list(protocol()) n_classes = len(np.asarray(base_bags[0][1])) self._classes_ = np.arange(n_classes) if val_protocol is not None: val_bags_list = list(val_protocol()) else: n_val = max(1, int(len(base_bags) * self.val_split)) val_bags_list = base_bags[:n_val] base_bags = base_bags[n_val:] rng = random.Random(self.random_state) n_mixed = round(len(base_bags) * mix_bags_proportion) if mix_bags else 0 def train_bags(): bags = list(base_bags) if n_mixed > 0: for _ in range(n_mixed): a, b = rng.choice(base_bags), rng.choice(base_bags) bags.append(mix_two_bags(a, b, self.bag_size, rng)) rng.shuffle(bags) return bags def val_bags(): return val_bags_list n_features = np.asarray(base_bags[0][0]).shape[1] self._fit_loop( train_bags, val_bags, n_features, n_bags_train=len(base_bags) + n_mixed, n_bags_val=len(val_bags_list) ) return self def _fit_loop(self, train_bags_fn, val_bags_fn, n_features, n_bags_train, n_bags_val): os.makedirs(self.checkpointdir, exist_ok=True) n_classes = len(self._classes_) fe = self.feature_extraction_module if fe is None: fe = IdentityFeatureExtractionModule(n_features) quantmodule = self._build_quantmodule(fe.output_size) self.model = BagNetworkModule( fe, quantmodule, n_classes, linear_sizes=self.linear_sizes, dropout=self.dropout, output_function=self.output_function ).to(self.device) optim = torch.optim.Adam(self.model.parameters(), lr=self.start_lr, weight_decay=self.weight_decay) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=self.patience, factor=self.lr_factor) early_stop = EarlyStop(self.patience, lower_is_better=True) best_state = copy.deepcopy(self.model.state_dict()) for epoch in range(self.train_epochs): self._run_epoch(train_bags_fn(), n_bags_train, optim, train=True, epoch=epoch) va_loss = self._run_epoch(val_bags_fn(), n_bags_val, optim=None, train=False, epoch=epoch) early_stop(va_loss, epoch) if early_stop.IMPROVED: best_state = copy.deepcopy(self.model.state_dict()) torch.save(best_state, self.checkpoint) elif early_stop.STOP: if self.verbose > 0: print(f'[{self._checkpoint_prefix}] training ended by patience exhausted at epoch {epoch}; ' f'restoring best model from epoch {early_stop.best_epoch}') break scheduler.step(va_loss) if optim.param_groups[0]['lr'] < self.end_lr: if self.verbose > 0: print(f'[{self._checkpoint_prefix}] early stopping in epoch {epoch} (learning rate below end_lr)') break self.model.load_state_dict(best_state) def _run_epoch(self, bags, n_bags, optim, train, epoch): self.model.train(mode=train) losses = [] pbar = tqdm(bags, total=n_bags, disable=self.verbose == 0) batch = [] def process_batch(batch): X, P = stack_bags(batch, self.device) if train: optim.zero_grad() P_hat = self.model.forward(X) loss = self.quant_loss(P, P_hat) + self._extra_loss() loss.backward() optim.step() else: with torch.no_grad(): P_hat = self.model.forward(X) loss = self.quant_loss(P, P_hat) return loss.item() for bag in pbar: batch.append(bag) if len(batch) == self.batch_size: losses.append(process_batch(batch)) batch = [] pbar.set_description( f'[{self._checkpoint_prefix}] epoch={epoch} {"train" if train else "val"}-' f'loss={np.mean(losses):.5f}' ) if batch: losses.append(process_batch(batch)) return np.mean(losses) if losses else float('inf') def predict(self, X): """ Generates a class prevalence estimate for the sample `X`, via a single forward pass of the trained network (the quantification module aggregates over however many instances are given, so `X` need not match the `bag_size` used during training). :param X: the test instances :return: `np.ndarray` of shape `(n_classes,)` with the class prevalence estimates """ self.model.eval() with torch.no_grad(): X_t = to_tensor(X, self.device).unsqueeze(0) prevalence = self.model.forward(X_t) return prevalence.cpu().numpy().flatten()