""" GMNet implementation. Ported from the reference implementation at https://github.com/pglez84/gmnet (the `GMNet`/ `DLQuantification` classes in that repo), adapted to QuaPy's own protocol-based sample generation (replacing that repo's custom, `quantificationlib`-backed bag generators), and reusing the shared bag-based training loop already factored out for :class:`quapy.method._histnet.HistNetQ` (see :class:`quapy.method._neural_bags.BagTrainedQuantifier`). The overall architecture is: one or more "GM branches" -- each a small per-branch feature extractor followed by a layer of Gaussian likelihoods (a :class:`_GMLayer`) evaluated at every instance of a bag -- concatenated and mean-pooled over the bag, followed by the shared quantification MLP head. Like HistNetQ (and QuaNet), GMNet is trained end-to-end by minimizing a quantification loss over samples ("bags") of known prevalence, rather than over individually labeled instances. Two deliberate deviations from the reference implementation, both required for the model to satisfy QuaPy's `predict(X)` contract (i.e., to be usable on a real test collection of arbitrary size, as opposed to only on bags resampled at the fixed `bag_size` used for training): * the original `GMNet_Module` reshapes each branch's per-instance likelihoods around a *fixed*, constructor-time `bag_size` (via `torch.nn.Unflatten(0, (-1, bag_size))`), which only works when every forward pass is fed bags of exactly that size. Here, the reshape is instead computed from the actual input shape at forward time (see :class:`_GMBranch`), which is equivalent when the bag size matches but also supports bags (or, at prediction time, whole test samples) of any other size. * the forward hooks used by the original code to capture each branch's pre-Gaussian latent activations (for the CKA regularization term) are replaced by simply storing that activation as an attribute during `forward` (see :attr:`_GMBranch.latent_activation`), since branches are now implemented with a plain `forward` method rather than an opaque `torch.nn.Sequential`. """ import numpy as np import scipy.spatial.distance import torch import torch.nn as nn import geotorch from quapy.method._neural_bags import BagTrainedQuantifier from quapy.protocol import UPP def _cka(latent_activations): """Feature-space linear CKA (Centered Kernel Alignment), averaged over every pair of latent activations, following the `CKARegularization` class in the reference implementation. Used to encourage the Gaussian components learned by different GM branches to capture complementary (dissimilar) aspects of the instances. :param latent_activations: a list of tensors, one per GM branch, all of shape (n_instances, dim_i) (dim_i may differ across branches) """ cka_sum = 0. n_pairs = 0 for i in range(len(latent_activations)): for j in range(i + 1, len(latent_activations)): x = latent_activations[i] y = latent_activations[j] x = x - torch.mean(x, dim=0, keepdim=True) y = y - torch.mean(y, dim=0, keepdim=True) dot_product_similarity = torch.norm(torch.matmul(x.t(), y)) ** 2 normalization_x = torch.norm(torch.matmul(x.t(), x)) normalization_y = torch.norm(torch.matmul(y.t(), y)) cka_sum = cka_sum + dot_product_similarity / (normalization_x * normalization_y) n_pairs += 1 return cka_sum / n_pairs class _GMLayer(nn.Module): """A layer of `num_gaussians` (unnormalized) Gaussian likelihoods, evaluated at every instance of a bag. `centers` and `covariance` are learned; `covariance` is constrained to stay positive-definite throughout training via `geotorch.positive_definite`. """ def __init__(self, n_features, num_gaussians): super().__init__() self.n_features = n_features self.num_gaussians = num_gaussians self.centers = nn.Parameter(torch.rand(num_gaussians, n_features)) self.covariance = nn.Parameter(torch.eye(n_features).repeat(num_gaussians, 1, 1)) geotorch.positive_definite(self, "covariance") # initialize the centers' covariance from the (squared, halved) nearest-neighbor distance # between the randomly initialized centers, so that gaussians start with a sensible spread centers = self.centers.detach().cpu().numpy() distances = scipy.spatial.distance.cdist(centers, centers) np.fill_diagonal(distances, np.inf) cov = (np.mean(np.min(distances, axis=1)) / 2) ** 2 self.covariance = torch.eye(n_features).repeat(num_gaussians, 1, 1) * cov def forward(self, x): # x: (batch_size, bag_size, n_features) centers = self.centers.unsqueeze(0).unsqueeze(0) # (1, 1, num_gaussians, n_features) diff = x.unsqueeze(2) - centers # (batch_size, bag_size, num_gaussians, n_features) cov_inv = torch.inverse(self.covariance) det_cov = torch.linalg.det(self.covariance) mahalanobis = torch.einsum('...i,...ij,...j->...', diff, cov_inv.unsqueeze(0).unsqueeze(0), diff) normalization_term = torch.log((2 * torch.pi) ** self.n_features * det_cov).unsqueeze(0).unsqueeze(0) log_probs = -0.5 * (mahalanobis + normalization_term) return torch.exp(log_probs) # (batch_size, bag_size, num_gaussians) class _GMBranch(nn.Module): """One GM branch: an optional small MLP mapping the (already feature-extracted) instances into a `gaussian_dimensions`-sized latent space, followed by a Sigmoid, a :class:`_GMLayer`, and a BatchNorm applied instance-wise (i.e., over the merged batch*bag_size dimension, matching the reference implementation). """ def __init__(self, input_size, num_gaussians, gaussian_dimensions, hidden_size_fe, dropout_fe): super().__init__() self.pre = nn.Sequential() prev_size = input_size latent_size = gaussian_dimensions if gaussian_dimensions is not None else input_size if gaussian_dimensions is not None: for j, layer_size in enumerate(hidden_size_fe or ()): self.pre.add_module(f'hidden_{j}', nn.Linear(prev_size, layer_size)) self.pre.add_module(f'leakyrelu_{j}', nn.LeakyReLU()) self.pre.add_module(f'dropout_{j}', nn.Dropout(dropout_fe)) prev_size = layer_size self.pre.add_module('latent_linear', nn.Linear(prev_size, gaussian_dimensions)) self.pre.add_module('sigmoid', nn.Sigmoid()) self.gm_layer = _GMLayer(n_features=latent_size, num_gaussians=num_gaussians) self.batch_norm = nn.BatchNorm1d(num_features=num_gaussians) self.output_size = num_gaussians self.latent_activation = None # populated on every forward(), read by GMNet's CKA regularization def forward(self, x): # x: (batch_size, bag_size, input_size) batch_size, bag_size = x.shape[0], x.shape[1] latent = self.pre(x) self.latent_activation = latent.reshape(-1, latent.shape[-1]) likelihoods = self.gm_layer(latent) # (batch_size, bag_size, num_gaussians) flat = self.batch_norm(likelihoods.reshape(batch_size * bag_size, -1)) return flat.reshape(batch_size, bag_size, -1) class _GMNetModule(nn.Module): """The quantification module for GMNet: one or more :class:`_GMBranch` instances, each producing a per-instance representation that is concatenated across branches and mean-pooled over the bag, as required by :class:`quapy.method._neural_bags.BagTrainedQuantifier`. """ def __init__(self, input_size, num_gaussians, n_gm_layers, gaussian_dimensions, hidden_size_fe=None, dropout_fe=0., cka_regularization=0.): super().__init__() if len(num_gaussians) != n_gm_layers: raise ValueError('num_gaussians should be a tuple of the same size as n_gm_layers') if len(gaussian_dimensions) != n_gm_layers: raise ValueError('gaussian_dimensions should be a tuple of the same size as n_gm_layers') self.n_gm_layers = n_gm_layers self.cka_regularization = cka_regularization self.branches = nn.ModuleList([ _GMBranch(input_size, num_gaussians[i], gaussian_dimensions[i], hidden_size_fe, dropout_fe) for i in range(n_gm_layers) ]) self.output_size = sum(num_gaussians) def forward(self, x): outputs = [branch(x) for branch in self.branches] return torch.mean(torch.cat(outputs, dim=-1), dim=1) def apply_regularization(self): """Whether the CKA regularization term should be added to the training loss: requires at least two GM branches (CKA is a pairwise measure) and a nonzero `cka_regularization` weight.""" return self.n_gm_layers > 1 and self.cka_regularization != 0 def regularization_term(self): latent_activations = [branch.latent_activation for branch in self.branches] return self.cka_regularization * _cka(latent_activations) class GMNet(BagTrainedQuantifier): """ Implementation of `GMNet `_, a neural network for quantification that represents each instance of a bag by its likelihood under one or more learned mixtures of Gaussians, mean-pools these representations over the bag, and predicts the class prevalence from the result, trained end-to-end by minimizing a quantification loss over many samples ("bags") of known prevalence. Like :class:`quapy.method._histnet.HistNetQ` and :class:`quapy.method.meta.QuaNet`, GMNet does not follow the classify-then-aggregate pattern of :class:`quapy.method.aggregative.AggregativeQuantifier`; it is instead trained and evaluated end-to-end on whole bags (see :class:`quapy.method._neural_bags.BagTrainedQuantifier` for the shared training/prediction logic, including the two entry points, :meth:`fit` and :meth:`fit_from_samples`). :param feature_extraction_module: a `torch.nn.Module` exposing an `output_size` attribute, used to embed each instance before it is passed to every GM branch. 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 n_gm_layers: number of GM branches (default 1). :param num_gaussians: number of gaussians per branch: either a single int (used for every branch) or a tuple/list of `n_gm_layers` ints (default 4). :param gaussian_dimensions: dimensionality of the latent space in which each branch's gaussians live: either a single int/None (used for every branch) or a tuple/list of `n_gm_layers` int/None values. If None for a given branch, that branch's gaussians operate directly on the feature-extracted instances, with no extra per-branch projection (default None). :param hidden_size_fe: sizes of the hidden layers of the small per-branch MLP that maps the feature-extracted instances into the latent space (only used when `gaussian_dimensions` is not None for the corresponding branch); default None (no hidden layers, i.e., a single linear projection). :param dropout_fe: dropout applied after each of the `hidden_size_fe` layers (default 0). :param cka_regularization: weight of the CKA regularization term encouraging the different branches' latent representations to be dissimilar; only applied when `n_gm_layers > 1` (default 0, i.e., disabled). :param linear_sizes: tuple of ints with the sizes of the linear layers used in the shared quantification head, after the GM branches (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, as well as for the random initialization of the GM branches (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 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, n_gm_layers=1, num_gaussians=4, gaussian_dimensions=None, hidden_size_fe=None, dropout_fe=0., cka_regularization=0., 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): super().__init__( feature_extraction_module=feature_extraction_module, linear_sizes=linear_sizes, dropout=dropout, output_function=output_function, bag_size=bag_size, n_bags_train=n_bags_train, n_bags_val=n_bags_val, train_epochs=train_epochs, patience=patience, start_lr=start_lr, end_lr=end_lr, lr_factor=lr_factor, weight_decay=weight_decay, quant_loss=quant_loss, batch_size=batch_size, protocol=protocol, protocol_params=protocol_params, val_split=val_split, device=device, random_state=random_state, checkpointdir=checkpointdir, checkpointname=checkpointname, verbose=verbose, ) self.n_gm_layers = n_gm_layers self.num_gaussians = num_gaussians if isinstance(num_gaussians, (tuple, list)) \ else [num_gaussians] * n_gm_layers self.gaussian_dimensions = gaussian_dimensions if isinstance(gaussian_dimensions, (tuple, list)) \ else [gaussian_dimensions] * n_gm_layers self.hidden_size_fe = hidden_size_fe self.dropout_fe = dropout_fe self.cka_regularization = cka_regularization @property def _checkpoint_prefix(self): return 'GMNet' def _build_quantmodule(self, n_features): torch.manual_seed(self.random_state) return _GMNetModule( input_size=n_features, num_gaussians=self.num_gaussians, n_gm_layers=self.n_gm_layers, gaussian_dimensions=self.gaussian_dimensions, hidden_size_fe=self.hidden_size_fe, dropout_fe=self.dropout_fe, cka_regularization=self.cka_regularization, ) def _extra_loss(self): quantmodule = self.model.quantmodule if quantmodule.apply_regularization(): return quantmodule.regularization_term() return 0.