Coverage for /opt/hostedtoolcache/Python/3.12.15/x64/lib/python3.12/site-packages/equine/equine_protonet.py: 96%
308 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-15 00:14 +0000
1# Copyright 2024, MASSACHUSETTS INSTITUTE OF TECHNOLOGY
2# Subject to FAR 52.227-11 – Patent Rights – Ownership by the Contractor (May 2014).
3# SPDX-License-Identifier: MIT
4from __future__ import annotations
6import io
7import warnings
8from collections import OrderedDict
9from collections.abc import Callable
10from datetime import datetime
11from enum import Enum
12from typing import Any, Optional
14import icontract
15import numpy as np
16import torch
17from beartype import beartype
18from scipy.stats import gaussian_kde
19from torch.utils.data import TensorDataset
20from tqdm import tqdm
22from .equine import Equine, EquineOutput
23from .utils import (
24 EQUINE_FORMAT_VERSION,
25 generate_episode,
26 generate_support,
27 generate_train_summary,
28 jit_archive_to_tensor,
29 load_checkpoint,
30 load_jit_archive,
31 mahalanobis_distance_nosq,
32 prepare_jit_module,
33 stratified_train_test_split,
34)
37#####################################
38class CovType(Enum):
39 """
40 Enum class for covariance types used in EQUINE.
41 """
43 UNIT = "unit"
44 DIAGONAL = "diag"
45 FULL = "full"
48PRED_COV_TYPE = CovType.DIAGONAL
49OOD_COV_TYPE = CovType.DIAGONAL
50DEFAULT_EPSILON = 1e-5
53def _kde_to_state(kde: gaussian_kde) -> dict[str, Any]:
54 """Reduce a fitted ``gaussian_kde`` to tensors/scalars for safe saving."""
55 return {
56 "dataset": torch.as_tensor(kde.dataset).clone(),
57 "bw_factor": float(kde.factor),
58 }
61def _kde_from_state(state: dict[str, Any]) -> gaussian_kde:
62 """Rebuild a ``gaussian_kde`` saved by ``_kde_to_state``."""
63 return gaussian_kde(state["dataset"].cpu().numpy(), bw_method=state["bw_factor"])
66COV_REG_TYPE = "epsilon"
69###############################################
72@beartype
73class Protonet(torch.nn.Module):
74 """
75 Private class that implements a prototypical neural network for use in EQUINE.
76 """
78 def __init__(
79 self,
80 embedding_model: torch.nn.Module,
81 emb_out_dim: int,
82 cov_type: CovType,
83 cov_reg_type: str,
84 epsilon: float,
85 device: str = "cpu",
86 ) -> None:
87 """
88 Protonet class constructor.
90 Parameters
91 ----------
92 embedding_model : torch.nn.Module
93 The PyTorch embedding model to generate logits with.
94 emb_out_dim : int
95 Dimension size of given embedding model's output.
96 cov_type : CovType
97 Type of covariance to use when computing distances [unit, diag, full].
98 cov_reg_type : str
99 Type of regularization to use when generating the covariance matrix [epsilon, shared].
100 epsilon : float
101 Epsilon value to use for covariance regularization.
102 device : str, optional
103 The device to train the protonet model on (defaults to cpu).
104 """
105 super().__init__()
106 self.embedding_model = embedding_model
107 self.cov_type = cov_type
108 self.cov_reg_type = cov_reg_type
109 self.epsilon = epsilon
110 self.emb_out_dim = emb_out_dim
111 self.to(device)
112 self.device = device
114 self.support: OrderedDict[int, torch.Tensor] = OrderedDict()
115 self.support_embeddings: OrderedDict[int, torch.Tensor] = OrderedDict()
116 self.model_head: torch.nn.Module = self.create_model_head(emb_out_dim)
117 self.model_head.to(device)
119 def create_model_head(self, emb_out_dim: int) -> torch.nn.Identity:
120 """
121 Method for adding a PyTorch layer on top of the given embedding model. This layer
122 is intended to offer extra degrees of freedom for distance learning in the embedding space.
124 Parameters
125 ----------
126 emb_out_dim : int
127 Dimension size of the embedding model output.
129 Returns
130 -------
131 torch.nn.Linear
132 The created PyTorch model layer.
133 """
134 return torch.nn.Identity() # (emb_out_dim, emb_out_dim)
136 def compute_embeddings(self, X: torch.Tensor) -> torch.Tensor:
137 """
138 Method for calculating model embeddings using both the given embedding model and the added model head.
140 Parameters
141 ----------
142 X : torch.Tensor
143 Input tensor to compute embeddings on.
145 Returns
146 -------
147 torch.Tensor
148 Fully computed embedding tensors for the given X tensor.
149 """
150 model_embeddings = self.embedding_model(X.to(self.device))
151 head_embeddings = self.model_head(model_embeddings)
152 return head_embeddings
154 @icontract.require(lambda self: len(self.support_embeddings) > 0)
155 def compute_prototypes(self) -> torch.Tensor:
156 """
157 Method for computing class prototypes based on given support examples.
158 ``Prototypes'' in this context are the means of the support embeddings for each class.
160 Returns
161 -------
162 torch.Tensor
163 Tensors of prototypes for each of the given classes in the support.
164 """
165 # Compute prototype for each class
166 proto_list = []
167 for label in self.support_embeddings: # look at doing functorch
168 class_prototype = torch.mean(self.support_embeddings[label], dim=0)
169 proto_list.append(class_prototype)
171 prototypes = torch.stack(proto_list)
173 return prototypes
175 @icontract.require(lambda self: len(self.support_embeddings) > 0)
176 def compute_covariance(self, cov_type: CovType) -> torch.Tensor:
177 """
178 Method for generating the (regularized) support example covariance matrix(es) used for calculating distances.
179 Note that this method is only called once per episode, and the resulting tensor is used for all queries.
181 Parameters
182 ----------
183 cov_type : CovType
184 Type of covariance to use [unit, diag, full].
186 Returns
187 -------
188 torch.Tensor
189 Tensor containing the generated regularized covariance matrix.
190 """
191 class_cov_dict = OrderedDict().fromkeys(
192 self.support_embeddings.keys(), torch.Tensor()
193 )
194 for label in self.support_embeddings.keys():
195 class_covariance = self.compute_covariance_by_type(
196 cov_type, self.support_embeddings[label]
197 )
198 class_cov_dict[label] = class_covariance
200 reg_covariance_dict = self.regularize_covariance(
201 class_cov_dict, cov_type, self.cov_reg_type
202 )
203 reg_covariance = torch.stack(list(reg_covariance_dict.values()))
205 return reg_covariance # TODO try putting everything on GPU with .to() and see if faster
207 def compute_covariance_by_type(
208 self, cov_type: CovType, embedding: torch.Tensor
209 ) -> torch.Tensor:
210 """
211 Select the appropriate covariance matrix type based on cov_type.
213 Parameters
214 ----------
215 cov_type : str
216 Type of covariance to use. Options are ['unit', 'diag', 'full'].
217 embedding : torch.Tensor
218 Embedding tensor to use when generating the covariance matrix.
220 Returns
221 -------
222 torch.Tensor
223 Tensor containing the requested covariance matrix.
224 """
225 if cov_type == CovType.FULL:
226 class_covariance = torch.cov(embedding.T)
227 elif cov_type == CovType.DIAGONAL:
228 class_covariance = torch.var(embedding, dim=0)
229 elif cov_type == CovType.UNIT:
230 class_covariance = torch.ones(self.emb_out_dim)
231 else:
232 raise ValueError
234 return class_covariance
236 def regularize_covariance(
237 self,
238 class_cov_dict: OrderedDict[int, torch.Tensor],
239 cov_type: CovType,
240 cov_reg_type: str,
241 ) -> OrderedDict[int, torch.Tensor]:
242 """
243 Method to add regularization to each class covariance matrix based on the selected regularization type.
245 Parameters
246 ----------
247 class_cov_dict : OrderedDict[int, torch.Tensor]
248 A dictionary containing each class and the corresponding covariance matrix.
249 cov_type : CovType
250 Type of covariance to use [unit, diag, full].
252 Returns
253 -------
254 dict[float, torch.Tensor]
255 Dictionary containing the regularized class covariance matrices.
256 """
258 if cov_type == CovType.FULL:
259 regularization = torch.diag(self.epsilon * torch.ones(self.emb_out_dim)).to(
260 self.device
261 )
262 elif cov_type == CovType.DIAGONAL:
263 regularization = self.epsilon * torch.ones(self.emb_out_dim).to(self.device)
264 elif cov_type == CovType.UNIT: 264 ↛ 267line 264 didn't jump to line 267 because the condition on line 264 was always true
265 regularization = torch.zeros(self.emb_out_dim).to(self.device)
267 if cov_reg_type == "shared":
268 if cov_type != CovType.FULL and cov_type != CovType.DIAGONAL: 268 ↛ 269line 268 didn't jump to line 269 because the condition on line 268 was never true
269 for label in self.support_embeddings:
270 class_cov_dict[label] = class_cov_dict[label] + regularization
271 warnings.warn(
272 "Covariance type UNIT is incompatible with shared regularization, \
273 reverting to epsilon regularization"
274 )
275 return class_cov_dict
277 shared_covariance = self.compute_shared_covariance(class_cov_dict, cov_type)
279 for label in self.support_embeddings:
280 num_class_support = self.support_embeddings[label].shape[0]
281 lamb = num_class_support / (num_class_support + 1)
283 class_cov_dict[label] = (
284 lamb * class_cov_dict[label]
285 + (1 - lamb) * shared_covariance
286 + regularization
287 )
289 elif cov_reg_type == "epsilon": 289 ↛ 295line 289 didn't jump to line 295 because the condition on line 289 was always true
290 for label in class_cov_dict.keys():
291 class_cov_dict[label] = (
292 class_cov_dict[label].to(self.device) + regularization
293 )
295 return class_cov_dict
297 def compute_shared_covariance(
298 self, class_cov_dict: OrderedDict[int, torch.Tensor], cov_type: CovType
299 ) -> torch.Tensor:
300 """
301 Method to calculate a shared covariance matrix.
303 The shared covariance matrix is calculated as the weighted average of the class covariance matrices,
304 where the weights are the number of support examples for each class. This is useful when the number of
305 support examples for each class is small.
307 Parameters
308 ----------
309 class_cov_dict : OrderedDict[int, torch.Tensor]
310 A dictionary containing each class and the corresponding covariance matrix.
311 cov_type : CovType
312 Type of covariance to use [unit, diag, full].
314 Returns
315 -------
316 torch.Tensor
317 Tensor containing the shared covariance matrix.
318 """
319 total_support = sum([x.shape[0] for x in class_cov_dict.values()])
321 if cov_type == CovType.FULL: 321 ↛ 322line 321 didn't jump to line 322 because the condition on line 321 was never true
322 shared_covariance = torch.zeros((self.emb_out_dim, self.emb_out_dim))
323 elif cov_type == CovType.DIAGONAL:
324 shared_covariance = torch.zeros(self.emb_out_dim)
325 else:
326 raise ValueError(
327 "Shared covariance can only be used with FULL or DIAGONAL (not UNIT) covariance types"
328 )
330 for label in class_cov_dict:
331 num_class_support = class_cov_dict[label].shape[0]
332 shared_covariance = (
333 shared_covariance + (num_class_support - 1) * class_cov_dict[label]
334 ) # undo N-1 div from cov
336 shared_covariance = shared_covariance / (
337 total_support - 1
338 ) # redo N-1 div for shared cov
340 return shared_covariance
342 @icontract.require(lambda X_embed, mu: X_embed.shape[-1] == mu.shape[-1])
343 @icontract.ensure(lambda result: torch.all(result >= 0))
344 def compute_distance(
345 self, X_embed: torch.Tensor, mu: torch.Tensor, cov: torch.Tensor
346 ) -> torch.Tensor:
347 """
348 Method to compute the distances to class prototypes for the given embeddings.
350 Parameters
351 ----------
352 X_embed : torch.Tensor
353 The embeddings of the query examples.
354 mu : torch.Tensor
355 The class prototypes (means of the support embeddings).
356 cov : torch.Tensor
357 The support covariance matrix.
359 Returns
360 -------
361 torch.Tensor
362 The calculated distances from each of the class prototypes for the given embeddings.
363 """
364 _queries = torch.unsqueeze(X_embed, 1) # examples x 1 x dimension
365 diff = torch.sub(mu, _queries)
367 if len(cov.shape) == 2: # (diagonal covariance)
368 # examples x classes x dimension
369 sq_diff = diff**2
370 div = torch.div(sq_diff.to(self.device), cov.to(self.device))
371 dist = torch.nan_to_num(div)
372 dist = torch.sum(dist, dim=2) # examples x classes
373 dist = dist.squeeze(dim=1)
374 dist = torch.sqrt(dist + self.epsilon) # examples x classes
375 else: # len(cov.shape) == 3: (full covariance)
376 diff = diff.permute(1, 2, 0) # classes x dimension x examples
377 dist = mahalanobis_distance_nosq(diff, cov)
378 dist = torch.sqrt(dist.permute(1, 0) + self.epsilon) # examples x classes
379 dist = dist.squeeze(dim=1)
380 return dist
382 def compute_classes(self, distances: torch.Tensor) -> torch.Tensor:
383 """
384 Method to compute predicted classes from distances via a softmax function.
386 Parameters
387 ----------
388 distances : torch.Tensor
389 The distances of embeddings to class prototypes.
391 Returns
392 -------
393 torch.Tensor
394 Tensor of class predictions.
395 """
396 softmax = torch.nn.functional.softmax(torch.neg(distances), dim=-1)
397 return softmax
399 def forward(self, X: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
400 """
401 Protonet forward function, generates class probability predictions and distances from prototypes.
403 Parameters
404 ----------
405 X : torch.Tensor
406 Input tensor of queries for generating predictions.
408 Returns
409 -------
410 tuple[torch.Tensor, torch.Tensor]
411 tuple containing class probability predictions, and class distances from prototypes.
412 """
413 if len(self.support) == 0 or len(self.support_embeddings) == 0:
414 raise ValueError(
415 "No support examples found. Protonet Model requires model support to \
416 be set with the 'update_support()' method before calling forward."
417 )
419 X_embed = self.compute_embeddings(X)
420 if X_embed.shape == torch.Size([self.emb_out_dim]):
421 X_embed = X_embed.unsqueeze(dim=0) # handle single examples
422 distances = self.compute_distance(X_embed, self.prototypes, self.covariance)
423 classes = self.compute_classes(distances)
425 return classes, distances
427 def update_support(self, support: OrderedDict[int, torch.Tensor]) -> None:
428 """
429 Method to update the support examples, and all the calculations that rely on them.
431 Parameters
432 ----------
433 support : OrderedDict
434 Ordered dict containing class labels and their associated support examples.
435 """
436 self.support = support # TODO torch.nn.ParameterDict(support)
438 support_embs = OrderedDict().fromkeys(support.keys(), torch.Tensor())
439 for label in support:
440 support_embs[label] = self.compute_embeddings(support[label])
442 self.support_embeddings = (
443 support_embs # TODO torch.nn.ParameterDict(support_embs)
444 )
446 self.prototypes: torch.Tensor = self.compute_prototypes()
448 if self.training is False:
449 self.compute_global_moments()
450 self.covariance: torch.Tensor = self.compute_covariance(
451 cov_type=PRED_COV_TYPE
452 )
453 else:
454 self.covariance: torch.Tensor = self.compute_covariance(
455 cov_type=self.cov_type
456 )
458 @icontract.require(lambda self: len(self.support_embeddings) > 0)
459 def compute_global_moments(self) -> None:
460 """Method to calculate the global moments of the support embeddings for use in OOD score generation"""
461 embeddings = torch.cat(list(self.support_embeddings.values()))
462 self.global_covariance = torch.unsqueeze(
463 self.compute_covariance_by_type(OOD_COV_TYPE, embeddings), dim=0
464 )
465 global_reg_input = OrderedDict().fromkeys([0], torch.Tensor())
466 global_reg_input[0] = self.global_covariance
467 self.global_covariance: torch.Tensor = self.regularize_covariance(
468 global_reg_input, OOD_COV_TYPE, "epsilon"
469 )[0]
470 self.global_mean: torch.Tensor = torch.mean(embeddings, dim=0)
473###############################################
474@beartype
475class EquineProtonet(Equine):
476 """
477 A class representing an EQUINE model that utilizes protonets and (optionally) relative Mahalanobis distances
478 to generate OOD and model confidence scores. This wraps any pytorch embedding neural network
479 and provides the `forward`, `predict`, `save`, and `load` methods required by Equine.
480 """
482 def __init__(
483 self,
484 embedding_model: torch.nn.Module,
485 emb_out_dim: int,
486 cov_type: CovType = CovType.UNIT,
487 relative_mahal: bool = True,
488 use_temperature: bool = False,
489 init_temperature: float = 1.0,
490 device: str = "cpu",
491 feature_names: Optional[list[str]] = None,
492 label_names: Optional[list[str]] = None,
493 ) -> None:
494 """
495 EquineProtonet class constructor
497 Parameters
498 ----------
499 embedding_model : torch.nn.Module
500 Neural Network feature embedding model.
501 emb_out_dim : int
502 The number of output features from the embedding model.
503 cov_type : CovType, optional
504 The type of covariance to use when training the protonet [UNIT, DIAG, FULL], by default CovType.UNIT.
505 relative_mahal : bool, optional
506 Use relative mahalanobis distance for OOD calculations. If false, uses standard mahalanobis distance instead, by default True.
507 use_temperature : bool, optional
508 Whether to use temperature scaling after training, by default False.
509 init_temperature : float, optional
510 What to use as the initial temperature (1.0 has no effect), by default 1.0.
511 device : str, optional
512 The device to train the equine model on (defaults to cpu).
513 feature_names : list[str], optional
514 List of strings of the names of the tabular features (ex ["duration", "fiat_mean", ...])
515 label_names : list[str], optional
516 List of strings of the names of the labels (ex ["streaming", "voip", ...])
517 """
518 super().__init__(
519 embedding_model,
520 device=device,
521 feature_names=feature_names,
522 label_names=label_names,
523 )
524 self.cov_type = cov_type
525 self.cov_reg_type = COV_REG_TYPE
526 self.relative_mahal = relative_mahal
527 self.emb_out_dim = emb_out_dim
528 self.epsilon = DEFAULT_EPSILON
529 self.outlier_score_kde: OrderedDict[int, gaussian_kde] = OrderedDict()
530 self.model_summary: dict[str, Any] = dict()
531 self.use_temperature = use_temperature
532 self.init_temperature = init_temperature
533 self.register_buffer(
534 "temperature", torch.Tensor(self.init_temperature * torch.ones(1))
535 )
537 self.model: torch.nn.Module = Protonet(
538 embedding_model,
539 self.emb_out_dim,
540 self.cov_type,
541 self.cov_reg_type,
542 self.epsilon,
543 device=device,
544 )
546 def forward(self, X: torch.Tensor) -> torch.Tensor:
547 """
548 Generates logits for classification based on the input tensor.
550 Parameters
551 ----------
552 X : torch.Tensor
553 The input tensor for generating predictions.
555 Returns
556 -------
557 torch.Tensor
558 The output class predictions.
559 """
560 preds, _ = self.model(X)
561 return preds
563 @icontract.require(lambda calib_frac: calib_frac > 0 and calib_frac < 1)
564 def train_model(
565 self,
566 dataset: TensorDataset,
567 num_episodes: int,
568 calib_frac: float = 0.2,
569 support_size: int = 25,
570 way: int = 3,
571 episode_size: int = 100,
572 loss_fn: Callable = torch.nn.functional.cross_entropy,
573 opt_class: Callable = torch.optim.Adam,
574 num_calibration_epochs: int = 2,
575 calibration_lr: float = 0.01,
576 ) -> dict[str, Any]:
577 """
578 Train or fine-tune an EquineProtonet model.
580 Parameters
581 ----------
582 dataset : TensorDataset
583 Input pytorch TensorDataset of training data for model.
584 num_episodes : int
585 The desired number of episodes to use for training.
586 calib_frac : float, optional
587 Fraction of given training data to reserve for model calibration, by default 0.2.
588 support_size : int, optional
589 Number of support examples to generate for each class, by default 25.
590 way : int, optional
591 Number of classes to train on per episode, by default 3.
592 episode_size : int, optional
593 Number of examples to use per episode, by default 100.
594 loss_fn : Callable, optional
595 A pytorch loss function, eg., torch.nn.CrossEntropyLoss(), by default torch.nn.functional.cross_entropy.
596 opt_class : Callable, optional
597 A pytorch optimizer, e.g., torch.optim.Adam, by default torch.optim.Adam.
598 num_calibration_epochs : int, optional
599 The desired number of epochs to use for temperature scaling, by default 2.
600 calibration_lr : float, optional
601 Learning rate for temperature scaling, by default 0.01.
603 Returns
604 -------
605 tuple[dict[str, Any], torch.Tensor, torch.Tensor]
606 A tuple containing the model summary, the held out calibration data, and the calibration labels.
607 """
608 self.train()
610 if self.use_temperature:
611 self.temperature: torch.Tensor = torch.Tensor(
612 self.init_temperature * torch.ones(1)
613 ).type_as(self.temperature)
615 X, Y = dataset[:]
617 self.validate_feature_label_names(X.shape[-1], torch.unique(Y).shape[0])
619 train_x, calib_x, train_y, calib_y = stratified_train_test_split(
620 X, Y, test_size=calib_frac
621 )
622 optimizer = opt_class(self.parameters())
624 train_x.to(self.device)
625 train_y.to(self.device)
626 calib_x.to(self.device)
627 calib_y.to(self.device)
629 for i in tqdm(range(num_episodes)):
630 optimizer.zero_grad()
632 support, episode_x, episode_y = generate_episode(
633 train_x, train_y, support_size, way, episode_size
634 )
635 self.model.update_support(support)
637 _, dists = self.model(episode_x)
638 loss_value = loss_fn(
639 torch.neg(dists).to(self.device), episode_y.to(self.device)
640 )
641 loss_value.backward()
642 optimizer.step()
644 self.eval()
645 full_support = generate_support(
646 train_x,
647 train_y,
648 support_size,
649 selected_labels=torch.unique(train_y).tolist(),
650 )
652 self.model.update_support(
653 full_support
654 ) # update support with final selected examples
656 X_embed = self.model.compute_embeddings(calib_x)
657 pred_probs, dists = self.model(calib_x)
658 ood_dists = self._compute_ood_dist(X_embed, pred_probs, dists)
659 self._fit_outlier_scores(ood_dists, calib_y)
661 if self.use_temperature:
662 self.calibrate_temperature(
663 calib_x, calib_y, num_calibration_epochs, calibration_lr
664 )
666 date_trained = datetime.now().strftime("%m/%d/%Y, %H:%M:%S")
667 self.train_summary: dict[str, Any] = generate_train_summary(
668 self, train_y, date_trained
669 )
670 return_dict: dict[str, Any] = dict()
671 return_dict["train_summary"] = self.train_summary
672 return_dict["calib_x"] = calib_x
673 return_dict["calib_y"] = calib_y
674 return return_dict
676 def calibrate_temperature(
677 self,
678 calib_x: torch.Tensor,
679 calib_y: torch.Tensor,
680 num_calibration_epochs: int = 1,
681 calibration_lr: float = 0.01,
682 ) -> None:
683 """
684 Fine-tune the temperature after training. Note that this function is also run at the conclusion of train_model.
686 Parameters
687 ----------
688 calib_x : torch.Tensor
689 Training data to be used for temperature calibration.
690 calib_y : torch.Tensor
691 Labels corresponding to `calib_x`.
692 num_calibration_epochs : int, optional
693 Number of epochs to tune temperature, by default 1.
694 calibration_lr : float, optional
695 Learning rate for temperature optimization, by default 0.01.
697 Returns
698 -------
699 None
700 """
701 self.temperature.requires_grad = True
702 optimizer = torch.optim.Adam([self.temperature], lr=calibration_lr)
703 for t in range(num_calibration_epochs):
704 optimizer.zero_grad()
705 with torch.no_grad():
706 pred_probs, dists = self.model(calib_x)
707 dists = dists.to(self.device) / self.temperature.to(self.device)
708 loss = torch.nn.functional.cross_entropy(
709 torch.neg(dists).to(self.device), calib_y.to(torch.long).to(self.device)
710 )
711 loss.backward()
712 optimizer.step()
713 self.temperature.requires_grad = False
715 @icontract.ensure(lambda self: len(self.model.support_embeddings) > 0)
716 def _fit_outlier_scores(
717 self, ood_dists: torch.Tensor, calib_y: torch.Tensor
718 ) -> None:
719 """
720 Private function to fit outlier scores with a kernel density estimate (KDE).
722 Parameters
723 ----------
724 ood_dists : torch.Tensor
725 Tensor of computed OOD distances.
726 calib_y : torch.Tensor
727 Tensor of class labels for `ood_dists` examples.
729 Returns
730 -------
731 None
732 """
733 for label in self.model.support_embeddings.keys():
734 class_ood_dists = ood_dists[calib_y == int(label)].cpu().detach().numpy()
735 class_kde = gaussian_kde(class_ood_dists) # TODO convert to torch func
736 self.outlier_score_kde[label] = class_kde
738 def _compute_outlier_scores(self, ood_dists, predictions) -> torch.Tensor:
739 """
740 Private function to compute OOD scores using the calculated kernel density estimate (KDE).
742 Parameters
743 ----------
744 ood_dists : torch.Tensor
745 Tensor of computed OOD distances.
746 predictions : torch.Tensor
747 Tensor of model protonet predictions.
749 Returns
750 -------
751 torch.Tensor
752 Tensor of OOD scores for the given examples.
753 """
754 ood_scores = torch.zeros_like(ood_dists)
755 for i in range(len(predictions)):
756 # Use KDE and RMD corresponding to the predicted class
757 predicted_class = int(torch.argmax(predictions[i, :]))
758 p_value = self.outlier_score_kde[int(predicted_class)].integrate_box_1d(
759 ood_dists[i].detach().cpu().numpy(), np.inf
760 )
761 ood_scores[i] = 1.0 - np.clip(p_value, 0.0, 1.0)
763 return ood_scores
765 @icontract.ensure(lambda result: len(result) > 0)
766 def _compute_ood_dist(
767 self,
768 X_embeddings: torch.Tensor,
769 predictions: torch.Tensor,
770 distances: torch.Tensor,
771 ) -> torch.Tensor:
772 """
773 Private function to compute OOD distances using a distance function.
775 Parameters
776 ----------
777 X_embeddings : torch.Tensor
778 Tensor of example embeddings.
779 predictions : torch.Tensor
780 Tensor of model protonet predictions for the given embeddings.
781 distances : torch.Tensor
782 Tensor of calculated protonet distances for the given embeddings.
784 Returns
785 -------
786 torch.Tensor
787 Tensor of OOD distances for the given embeddings.
788 """
789 preds = torch.argmax(predictions, dim=1)
790 preds = preds.unsqueeze(dim=-1)
791 # Calculate (Relative) Mahalanobis Distance:
792 if self.relative_mahal:
793 null_distance = self.model.compute_distance(
794 X_embeddings, self.model.global_mean, self.model.global_covariance
795 )
796 null_distance = null_distance.unsqueeze(dim=-1)
797 ood_dist = distances.gather(1, preds) - null_distance
798 else:
799 ood_dist = distances.gather(1, preds)
801 ood_dist = torch.reshape(ood_dist, (-1,))
802 return ood_dist
804 def predict(self, X: torch.Tensor) -> EquineOutput:
805 """Predict function for EquineProtonet, inherited and implemented from Equine.
807 Parameters
808 ----------
809 X : torch.Tensor
810 Input tensor.
812 Returns
813 -------
814 EquineOutput
815 Output object containing prediction probabilities and OOD scores.
816 """
817 X_embed = self.model.compute_embeddings(X)
818 if X_embed.shape == torch.Size([self.model.emb_out_dim]):
819 X_embed = X_embed.unsqueeze(dim=0) # Handle single examples
820 preds, dists = self.model(X)
821 if self.use_temperature:
822 dists = dists / self.temperature
823 preds = torch.softmax(torch.negative(dists), dim=1)
824 ood_dist = self._compute_ood_dist(X_embed, preds, dists)
825 ood_scores = self._compute_outlier_scores(ood_dist, preds)
827 self.validate_feature_label_names(X.shape[-1], preds.shape[-1])
829 return EquineOutput(classes=preds, ood_scores=ood_scores, embeddings=X_embed)
831 @icontract.require(lambda calib_frac: (calib_frac > 0.0) and (calib_frac < 1.0))
832 def update_support(
833 self,
834 support_x: torch.Tensor,
835 support_y: torch.Tensor,
836 calib_frac: float,
837 label_names: Optional[list[str]] = None,
838 ) -> None:
839 """Function to update protonet support examples with given examples.
841 Parameters
842 ----------
843 support_x : torch.Tensor
844 Tensor containing support examples for protonet.
845 support_y : torch.Tensor
846 Tensor containing labels for given support examples.
847 calib_frac : float
848 Fraction of given support data to use for OOD calibration.
849 label_names : list[str], optional
850 List of strings of the names of the labels (ex ["streaming", "voip", ...])
852 Returns
853 -------
854 None
855 """
857 support_x, calib_x, support_y, calib_y = stratified_train_test_split(
858 support_x, support_y, test_size=calib_frac
859 )
860 labels, counts = torch.unique(support_y, return_counts=True)
861 if label_names is not None: 861 ↛ 862line 861 didn't jump to line 862 because the condition on line 861 was never true
862 self.label_names = label_names
863 self.validate_feature_label_names(support_x.shape[-1], labels.shape[0])
865 support = OrderedDict()
866 for label, count in list(zip(labels.tolist(), counts.tolist())):
867 class_support = generate_support(
868 support_x,
869 support_y,
870 support_size=count,
871 selected_labels=[label],
872 )
873 support.update(class_support)
875 self.model.update_support(support)
877 X_embed = self.model.compute_embeddings(calib_x)
878 preds, dists = self.model(calib_x)
879 ood_dists = self._compute_ood_dist(X_embed, preds, dists)
881 self._fit_outlier_scores(ood_dists, calib_y)
883 @icontract.require(lambda self: len(self.model.support) > 0)
884 def get_support(self) -> OrderedDict[int, torch.Tensor]:
885 """
886 Get the support examples for the model.
888 Returns
889 -------
890 OrderedDict[int, torch.Tensor]
891 The support examples for the model.
892 """
893 return self.model.support
895 @icontract.require(lambda self: len(self.model.prototypes) > 0)
896 def get_prototypes(self) -> torch.Tensor:
897 """
898 Get the prototypes for the model (the class means of the support embeddings).
900 Returns
901 -------
902 torch.Tensor
903 The prototpes for the model.
904 """
905 return self.model.prototypes
907 def save(self, path: str) -> None:
908 """
909 Save all model parameters to a file.
911 Parameters
912 ----------
913 path : str
914 Filename to write the model.
916 Returns
917 -------
918 None
919 """
920 model_settings = {
921 "cov_type": self.cov_type.value,
922 "emb_out_dim": self.emb_out_dim,
923 "use_temperature": self.use_temperature,
924 "init_temperature": self.temperature.item(),
925 "relative_mahal": self.relative_mahal,
926 "device": self.device,
927 }
929 jit_model = torch.jit.script(prepare_jit_module(self.model.embedding_model))
930 buffer = io.BytesIO()
931 torch.jit.save(jit_model, buffer)
933 # Plain values only, so load_checkpoint can read it with weights_only=True.
934 save_data = {
935 "equine_format_version": EQUINE_FORMAT_VERSION,
936 "embed_jit_save": jit_archive_to_tensor(buffer),
937 "feature_names": self.feature_names,
938 "label_names": self.label_names,
939 "model_head_save": self.model.model_head.state_dict(),
940 "outlier_kde": {
941 int(label): _kde_to_state(kde)
942 for label, kde in self.outlier_score_kde.items()
943 },
944 "settings": model_settings,
945 "support": {int(label): x for label, x in self.model.support.items()},
946 "train_summary": self.train_summary,
947 }
949 torch.save(save_data, path) # TODO allow model checkpointing
951 @classmethod
952 def load(cls, path: str, device: Optional[str] = None) -> Equine:
953 """
954 Load a previously saved EquineProtonet model.
956 Files from earlier EQUINE releases must first be rewritten with
957 ``python -m equine.upgrade``.
959 Parameters
960 ----------
961 path : str
962 The filename of the saved model.
963 device : Optional[str]
964 The device to load the model onto.
966 Returns
967 -------
968 EquineProtonet
969 The reconstituted EquineProtonet object.
971 Raises
972 ------
973 ValueError
974 If the file is not in the current EQUINE format.
975 """
976 model_save = load_checkpoint(path, map_location=device)
977 return cls._from_checkpoint(model_save, device)
979 @classmethod
980 def _from_checkpoint(
981 cls, model_save: dict[str, Any], device: Optional[str] = None
982 ) -> Equine:
983 """
984 Rebuild an EquineProtonet from an already-loaded checkpoint dictionary.
986 Parameters
987 ----------
988 model_save : dict[str, Any]
989 The dictionary returned by ``utils.load_checkpoint``.
990 device : Optional[str]
991 Device override for the reconstituted model.
993 Returns
994 -------
995 EquineProtonet
996 The reconstituted EquineProtonet object.
997 """
998 support = OrderedDict(
999 (int(label), x) for label, x in model_save.get("support").items()
1000 )
1002 # Explicitly pass map_location for the jit_model as well
1003 jit_model = load_jit_archive(model_save.get("embed_jit_save"), device)
1005 settings = dict(model_save.get("settings"))
1006 settings["cov_type"] = CovType(settings["cov_type"])
1007 # Allow the user to override the saved device state dynamically
1008 if device is not None: 1008 ↛ 1009line 1008 didn't jump to line 1009 because the condition on line 1008 was never true
1009 settings["device"] = device
1011 eq_model = cls(jit_model, **settings)
1013 eq_model.model.model_head.load_state_dict(model_save.get("model_head_save"))
1014 eq_model.eval()
1015 eq_model.model.update_support(support)
1017 eq_model.feature_names = model_save.get("feature_names")
1018 eq_model.label_names = model_save.get("label_names")
1019 eq_model.outlier_score_kde = OrderedDict(
1020 (int(label), _kde_from_state(kde))
1021 for label, kde in model_save.get("outlier_kde").items()
1022 )
1023 eq_model.train_summary = model_save.get("train_summary")
1025 return eq_model