Coverage for /opt/hostedtoolcache/Python/3.12.15/x64/lib/python3.12/site-packages/equine/equine_gp.py: 96%
280 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
5import io
6import math
7from collections import OrderedDict
8from collections.abc import Callable, Iterable
9from datetime import datetime
10from typing import Any, Optional, Union
12import icontract
13import torch
14from beartype import beartype
15from torch.utils.data import DataLoader, Dataset
16from torchmetrics.metric import Metric
17from tqdm import tqdm
19from .equine import Equine, EquineOutput
20from .utils import (
21 EQUINE_FORMAT_VERSION,
22 generate_support,
23 generate_train_summary,
24 jit_archive_to_tensor,
25 load_checkpoint,
26 load_jit_archive,
27 prepare_jit_module,
28)
30BatchType = tuple[torch.Tensor, ...]
31# -------------------------------------------------------------------------------
32# Note that the below code for
33# * `_random_ortho`,
34# * `_RandomFourierFeatures``, and
35# * `_Laplace`
36# is copied and modified from https://github.com/y0ast/DUE/blob/main/due/sngp.py
37# under its original MIT license, redisplayed here:
38# -------------------------------------------------------------------------------
39# MIT License
40#
41# Copyright (c) 2021 Joost van Amersfoort
42#
43# Permission is hereby granted, free of charge, to any person obtaining a copy
44# of this software and associated documentation files (the "Software"), to deal
45# in the Software without restriction, including without limitation the rights
46# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
47# copies of the Software, and to permit persons to whom the Software is
48# furnished to do so, subject to the following conditions:
49#
50# The above copyright notice and this permission notice shall be included in all
51# copies or substantial portions of the Software.
52#
53# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
54# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
55# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
56# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
57# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
58# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
59# SOFTWARE.
60# ------------------------------------------------------------------------------
61# Following the recommendation of their README at https://github.com/y0ast/DUE
62# we encourage anyone using this code in their research to cite the following papers:
63#
64# @article{van2021on,
65# title={On Feature Collapse and Deep Kernel Learning for Single Forward Pass Uncertainty},
66# author={van Amersfoort, Joost and Smith, Lewis and Jesson, Andrew and Key, Oscar and Gal, Yarin},
67# journal={arXiv preprint arXiv:2102.11409},
68# year={2021}
69# }
70#
71# @article{liu2020simple,
72# title={Simple and principled uncertainty estimation with deterministic deep learning via distance awareness},
73# author={Liu, Jeremiah and Lin, Zi and Padhy, Shreyas and Tran, Dustin and Bedrax Weiss, Tania and Lakshminarayanan, Balaji},
74# journal={Advances in Neural Information Processing Systems},
75# volume={33},
76# pages={7498--7512},
77# year={2020}
78# }
81@beartype
82def _random_ortho(n: int, m: int) -> torch.Tensor:
83 """
84 Generate a random orthonormal matrix.
86 Parameters
87 ----------
88 n : int
89 The number of rows.
90 m : int
91 The number of columns.
93 Returns
94 -------
95 torch.Tensor
96 The random orthonormal matrix.
97 """
98 q, _ = torch.linalg.qr(torch.randn(n, m))
99 return q
102@beartype
103class _RandomFourierFeatures(torch.nn.Module):
104 """
105 A private class to generate random Fourier features for the embedding model.
106 """
108 def __init__(
109 self, in_dim: int, num_random_features: int, feature_scale: Optional[float]
110 ) -> None:
111 """
112 Initialize the _RandomFourierFeatures module, which generates random Fourier features
113 for the embedding model.
115 Parameters
116 ----------
117 in_dim : int
118 The input dimensionality.
119 num_random_features : int
120 The number of random Fourier features to generate.
121 feature_scale : Optional[float]
122 The scaling factor for the random Fourier features. If None, defaults to sqrt(num_random_features / 2).
123 """
124 super().__init__()
125 if feature_scale is None: 125 ↛ 126line 125 didn't jump to line 126 because the condition on line 125 was never true
126 feature_scale = math.sqrt(num_random_features / 2)
128 self.register_buffer("feature_scale", torch.tensor(feature_scale))
130 if num_random_features <= in_dim: 130 ↛ 131line 130 didn't jump to line 131 because the condition on line 130 was never true
131 W: torch.Tensor = _random_ortho(in_dim, num_random_features)
132 else:
133 # generate blocks of orthonormal rows which are not necessarily orthonormal
134 # to each other.
135 dim_left = num_random_features
136 ws = []
137 while dim_left > in_dim:
138 ws.append(_random_ortho(in_dim, in_dim))
139 dim_left -= in_dim
140 ws.append(_random_ortho(in_dim, dim_left))
141 W: torch.Tensor = torch.cat(ws, 1)
143 feature_norm = torch.randn(W.shape) ** 2
145 W = W * feature_norm.sum(0).sqrt()
146 self.register_buffer("W", W)
148 b: torch.Tensor = torch.empty(num_random_features).uniform_(0, 2 * math.pi)
149 self.register_buffer("b", b)
151 def forward(self, x: torch.Tensor) -> torch.Tensor:
152 """
153 Compute the forward pass of the _RandomFourierFeatures module.
155 Parameters
156 ----------
157 x : torch.Tensor
158 The input tensor of shape (batch_size, in_dim).
160 Returns
161 -------
162 torch.Tensor
163 The output tensor of shape (batch_size, num_random_features).
164 """
165 k = torch.cos(x @ self.W + self.b)
166 k = k / self.feature_scale
168 return k
171class _Laplace(torch.nn.Module):
172 """
173 A private class to compute a Laplace approximation to a Gaussian Process (GP)
174 """
176 def __init__(
177 self,
178 feature_extractor: torch.nn.Module,
179 num_deep_features: int,
180 num_gp_features: int,
181 normalize_gp_features: bool,
182 num_random_features: int,
183 num_outputs: int,
184 feature_scale: Optional[float],
185 mean_field_factor: Optional[float], # required for classification problems
186 ridge_penalty: float = 1.0,
187 ) -> None:
188 """
189 Initialize the _Laplace module.
191 Parameters
192 ----------
193 feature_extractor : torch.nn.Module
194 The feature extractor module.
195 num_deep_features : int
196 The number of features output by the feature extractor.
197 num_gp_features : int
198 The number of features to use in the Gaussian process.
199 normalize_gp_features : bool
200 Whether to normalize the GP features.
201 num_random_features : int
202 The number of random Fourier features to use.
203 num_outputs : int
204 The number of outputs of the model.
205 feature_scale : Optional[float]
206 The scaling factor for the random Fourier features.
207 mean_field_factor : Optional[float]
208 The mean-field factor for the Gaussian-Softmax approximation.
209 Required for classification problems.
210 ridge_penalty : float, optional
211 The ridge penalty for the Laplace approximation.
212 """
213 super().__init__()
214 self.feature_extractor = feature_extractor
215 self.mean_field_factor = mean_field_factor
216 self.ridge_penalty = ridge_penalty
217 self.train_batch_size = 0 # to be set later
219 if num_gp_features > 0: 219 ↛ 229line 219 didn't jump to line 229 because the condition on line 219 was always true
220 self.num_gp_features = num_gp_features
221 random_matrix: torch.Tensor = torch.normal(
222 0, 0.05, (num_gp_features, num_deep_features)
223 )
224 self.register_buffer("random_matrix", random_matrix)
225 self.jl: Callable = lambda x: torch.nn.functional.linear(
226 x, self.random_matrix
227 )
228 else:
229 self.num_gp_features: int = num_deep_features
230 self.jl: Callable = lambda x: x # Identity
232 self.normalize_gp_features = normalize_gp_features
233 if normalize_gp_features: 233 ↛ 236line 233 didn't jump to line 236 because the condition on line 233 was always true
234 self.normalize: torch.nn.LayerNorm = torch.nn.LayerNorm(num_gp_features)
236 self.rff: _RandomFourierFeatures = _RandomFourierFeatures(
237 num_gp_features, num_random_features, feature_scale
238 )
239 self.beta: torch.nn.Linear = torch.nn.Linear(num_random_features, num_outputs)
241 self.num_data = 0 # to be set later
242 self.register_buffer("seen_data", torch.tensor(0))
244 precision = torch.eye(num_random_features) * self.ridge_penalty
245 self.register_buffer("precision", precision)
247 self.recompute_covariance = True
248 self.register_buffer("covariance", torch.eye(num_random_features))
249 self.training_parameters_set = False
251 def reset_precision_matrix(self) -> None:
252 """
253 Reset the precision matrix to the identity matrix times the ridge penalty.
254 """
255 identity = torch.eye(self.precision.shape[0], device=self.precision.device)
256 self.precision: torch.Tensor = identity * self.ridge_penalty
257 self.seen_data: torch.Tensor = torch.tensor(0)
258 self.recompute_covariance = True
260 @icontract.require(lambda num_data: num_data > 0)
261 @icontract.require(
262 lambda num_data, batch_size: (0 < batch_size) & (batch_size <= num_data)
263 )
264 def set_training_params(self, num_data: int, batch_size: int) -> None:
265 """
266 Set the training parameters for the Laplace approximation.
268 Parameters
269 ----------
270 num_data : int
271 The total number of data points.
272 batch_size : int
273 The batch size to use during training.
274 """
275 self.num_data: int = num_data
276 self.train_batch_size: int = batch_size
277 self.training_parameters_set: bool = True
279 @icontract.require(lambda mean_field_factor: mean_field_factor is not None)
280 def mean_field_logits(
281 self, logits: torch.Tensor, pred_cov: torch.Tensor, mean_field_factor: float
282 ) -> torch.Tensor:
283 """
284 Compute the mean-field logits for the Gaussian-Softmax approximation.
286 Parameters
287 ----------
288 logits : torch.Tensor
289 The logits tensor of shape (batch_size, num_outputs).
290 pred_cov : torch.Tensor
291 The predicted covariance matrix of shape (batch_size, batch_size).
292 mean_field_factor : float
293 Diagonal scaling factor
295 Returns
296 -------
297 torch.Tensor
298 The mean-field logits tensor of shape (batch_size, num_outputs).
299 """
300 # Mean-Field approximation as alternative to MC integration of Gaussian-Softmax
301 # Based on: https://arxiv.org/abs/2006.07584
303 logits_scale = torch.sqrt(1.0 + torch.diag(pred_cov) * mean_field_factor)
304 if mean_field_factor > 0: 304 ↛ 307line 304 didn't jump to line 307 because the condition on line 304 was always true
305 logits = logits / logits_scale.unsqueeze(-1)
307 return logits
309 @icontract.require(lambda self: self.training_parameters_set)
310 def forward(
311 self, x: torch.Tensor
312 ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
313 """
314 Compute the forward pass of the Laplace approximation to the Gaussian Process.
316 Parameters
317 ----------
318 x : torch.Tensor
319 The input tensor of shape (batch_size, num_features).
321 Returns
322 -------
323 Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]
324 If the model is in training mode, returns the predicted mean of shape (batch_size, 1).
325 If the model is in evaluation mode, returns a tuple containing the predicted mean of shape (batch_size, 1)
326 and the predicted covariance matrix of shape (batch_size, batch_size).
327 """
328 f = self.feature_extractor(x)
329 f_reduc = self.jl(f)
330 if self.normalize_gp_features: 330 ↛ 333line 330 didn't jump to line 333 because the condition on line 330 was always true
331 f_reduc = self.normalize(f_reduc)
333 k = self.rff(f_reduc)
335 pred = self.beta(k)
337 if self.training:
338 precision_minibatch = k.t() @ k
339 self.precision += precision_minibatch
340 self.seen_data += x.shape[0]
342 assert self.seen_data <= self.num_data, (
343 "Did not reset precision matrix at start of epoch"
344 )
345 else:
346 assert self.seen_data > (self.num_data - self.train_batch_size), (
347 "Not seen sufficient data for precision matrix"
348 )
350 if self.recompute_covariance:
351 with torch.no_grad():
352 eps = 1e-7
353 jitter = eps * torch.eye(
354 self.precision.shape[1],
355 device=self.precision.device,
356 )
357 u, info = torch.linalg.cholesky_ex(self.precision + jitter)
358 assert (info == 0).all(), "Precision matrix inversion failed!"
359 torch.cholesky_inverse(u, out=self.covariance)
361 self.recompute_covariance: bool = False
363 with torch.no_grad():
364 pred_cov = k @ ((self.covariance @ k.t()) * self.ridge_penalty)
366 if self.mean_field_factor is None: 366 ↛ 367line 366 didn't jump to line 367 because the condition on line 366 was never true
367 return pred, pred_cov
368 else:
369 pred = self.mean_field_logits(pred, pred_cov, self.mean_field_factor)
371 return pred
374# -------------------------------------------------------------------------------
375# EquineGP, below, demonstrates how to adapt that approach in EQUINE
376@beartype
377class EquineGP(Equine):
378 """
379 An example of an EQUINE model that builds upon the approach in "Spectral Norm
380 Gaussian Processes" (SNGP). This wraps any pytorch embedding neural network and provides
381 the `forward`, `predict`, `save`, and `load` methods required by Equine.
383 Notes
384 -----
385 Although this model build upon the approach in SNGP, it does not enforce the spectral normalization
386 and ResNet architecture required for SNGP. Instead, it is a simple wrapper around
387 any pytorch embedding neural network. Your mileage may vary.
388 """
390 def __init__(
391 self,
392 embedding_model: torch.nn.Module,
393 emb_out_dim: int,
394 num_classes: int,
395 num_random_features: int = 1024,
396 init_temperature: float = 1.0,
397 device: str = "cpu",
398 feature_names: Optional[list[str]] = None,
399 label_names: Optional[list[str]] = None,
400 ) -> None:
401 """
402 Initialize the EquineGP model.
404 Parameters
405 ----------
406 embedding_model : torch.nn.Module
407 Neural Network feature embedding.
408 emb_out_dim : int
409 The number of deep features from the feature embedding.
410 num_classes : int
411 The number of output classes this model predicts.
412 num_random_features : int
413 The dimension of the output of the RandomFourierFeatures operation
414 init_temperature : float, optional
415 What to use as the initial temperature (1.0 has no effect).
416 device : str, optional
417 Either 'cuda' or 'cpu'.
418 feature_names : list[str], optional
419 List of strings of the names of the tabular features (ex ["duration", "fiat_mean", ...])
420 label_names : list[str], optional
421 List of strings of the names of the labels (ex ["streaming", "voip", ...])
422 """
423 super().__init__(
424 embedding_model, feature_names=feature_names, label_names=label_names
425 )
426 self.num_deep_features = emb_out_dim
427 self.num_gp_features = emb_out_dim
428 self.normalize_gp_features = True
429 self.num_random_features = num_random_features
430 self.num_outputs = num_classes
431 self.mean_field_factor = 25
432 self.ridge_penalty = 1
433 self.feature_scale: float = 2.0
434 self.init_temperature = init_temperature
435 self.register_buffer(
436 "temperature", torch.Tensor(self.init_temperature * torch.ones(1))
437 )
438 self.model: _Laplace = _Laplace(
439 self.embedding_model,
440 self.num_deep_features,
441 self.num_gp_features,
442 self.normalize_gp_features,
443 self.num_random_features,
444 self.num_outputs,
445 self.feature_scale,
446 self.mean_field_factor,
447 self.ridge_penalty,
448 )
449 self.device_type = device
450 self.device: torch.device = torch.device(self.device_type)
451 self.model.to(self.device)
453 def train_model(
454 self,
455 dataset: Dataset,
456 loss_fn: Callable,
457 opt: torch.optim.Optimizer,
458 num_epochs: int,
459 scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
460 batch_size: int = 64,
461 validation_dataset: Optional[Dataset] = None,
462 val_metrics: Optional[Iterable[Metric]] = None,
463 vis_support: bool = False,
464 support_size: int = 25,
465 ) -> dict[str, Any]:
466 """
467 Train or fine-tune an EquineGP model.
469 Parameters
470 ----------
471 dataset : TensorDataset
472 An iterable, pytorch TensorDataset.
473 loss_fn : Callable
474 A pytorch loss function, e.g., torch.nn.CrossEntropyLoss().
475 opt : torch.optim.Optimizer
476 A pytorch optimizer, e.g., torch.optim.Adam().
477 num_epochs : int
478 The desired number of epochs to use for training.
479 scheduler : torch.optim.LRScheduler
480 A pytorch scheduler, if one is desired
481 validation_dataset: Dataset
482 If provided, will compute validation metrics on this dataset after each epoch of training
483 batch_size : int, optional
484 The number of samples to use per batch.
486 Returns
487 -------
488 dict[str, Any]
489 A dict containing a dict of summary stats and a dataloader for the calibration data.
491 """
493 self.validate_feature_label_names(dataset[0][0].shape[-1], self.num_outputs)
495 train_loader = DataLoader(
496 dataset, batch_size=batch_size, shuffle=True, drop_last=False
497 )
499 val_loader: Optional[DataLoader] = None
500 if validation_dataset is not None:
501 val_loader = DataLoader(
502 validation_dataset,
503 batch_size=batch_size,
504 shuffle=False,
505 drop_last=False,
506 )
508 self.model.set_training_params(len(dataset), batch_size)
509 val_metrics_outputs: Optional[list[list[float]]] = None
511 if validation_dataset is not None and val_metrics is not None:
512 val_metrics_outputs = [[] for i in range(len(list(val_metrics)))]
514 for _ in tqdm(range(num_epochs)):
515 self.model.train()
516 self.model.reset_precision_matrix()
517 epoch_loss = 0.0
518 for i, (xs, labels) in enumerate(train_loader):
519 opt.zero_grad()
520 xs = xs.to(self.device)
521 labels = labels.to(self.device)
522 yhats = self.model(xs)
523 loss = loss_fn(yhats, labels.to(torch.long))
524 loss.backward()
525 opt.step()
526 epoch_loss += loss.item()
527 if scheduler is not None:
528 scheduler.step()
529 self.model.eval()
530 # compute the validation metrics
531 if (
532 validation_dataset is not None
533 and val_loader is not None
534 and val_metrics is not None
535 and val_metrics_outputs is not None
536 ):
537 for _, (xs_val, labels_val) in enumerate(val_loader):
538 xs_val = xs_val.to(self.device)
539 labels_val = labels_val.to(self.device)
540 yhats_val = self.model(xs_val)
541 for metric in val_metrics:
542 metric.update(yhats_val, labels_val)
543 for i, metric in enumerate(val_metrics):
544 val_metrics_outputs[i].append(metric.compute())
545 if vis_support:
546 self.update_support(dataset.tensors[0], dataset.tensors[1], support_size)
548 _, train_y = dataset[:]
549 date_trained = datetime.now().strftime("%m/%d/%Y, %H:%M:%S")
550 self.train_summary: dict[str, Any] = generate_train_summary(
551 self, train_y, date_trained
552 )
554 return_dict: dict[str, Any] = dict()
555 return_dict["train_summary"] = self.train_summary
556 if validation_dataset is not None:
557 return_dict["val_metrics"] = val_metrics_outputs
559 return return_dict
561 def update_support(
562 self, support_x: torch.Tensor, support_y: torch.Tensor, support_size: int
563 ) -> None:
564 """Function to update protonet support examples with given examples.
566 Parameters
567 ----------
568 support_x : torch.Tensor
569 Tensor containing support examples for protonet.
570 support_y : torch.Tensor
571 Tensor containing labels for given support examples.
573 Returns
574 -------
575 None
576 """
578 labels, counts = torch.unique(support_y, return_counts=True)
579 support = OrderedDict()
580 for label, count in list(zip(labels.tolist(), counts.tolist())):
581 class_support = generate_support(
582 support_x,
583 support_y,
584 support_size=min(count, support_size),
585 selected_labels=[label],
586 )
587 support.update(class_support)
589 self.support = support
591 support_embeddings = OrderedDict().fromkeys(self.support.keys(), torch.Tensor())
592 for label in support:
593 support_embeddings[label] = self.compute_embeddings(support[label])
595 self.support_embeddings = support_embeddings
596 self.prototypes: torch.Tensor = self.compute_prototypes()
598 def compute_embeddings(self, x: torch.Tensor) -> torch.Tensor:
599 """
600 Method for computing deep embeddings for given input tensor.
602 Parameters
603 ----------
604 x : torch.Tensor
605 Input tensor for generating embeddings.
607 Returns
608 -------
609 torch.Tensor
610 Output embeddings .
611 """
612 f = self.model.feature_extractor(x)
613 f_reduc = self.model.jl(f)
614 if self.model.normalize_gp_features: 614 ↛ 617line 614 didn't jump to line 617 because the condition on line 614 was always true
615 f_reduc = self.model.normalize(f_reduc)
617 return self.model.rff(f_reduc)
619 @icontract.require(lambda self: len(self.support) > 0)
620 def compute_prototypes(self) -> torch.Tensor:
621 """
622 Method for computing class prototypes based on given support examples.
623 ``Prototypes'' in this context are the means of the support embeddings for each class.
625 Returns
626 -------
627 torch.Tensor
628 Tensors of prototypes for each of the given classes in the support.
629 """
630 # Compute support embeddings
631 support_embeddings = OrderedDict().fromkeys(self.support.keys())
632 for label in self.support:
633 support_embeddings[label] = self.compute_embeddings(self.support[label])
635 # Compute prototype for each class
636 proto_list = []
637 for label in self.support: # look at doing functorch
638 class_prototype = torch.mean(support_embeddings[label], dim=0) # type: ignore
639 proto_list.append(class_prototype)
641 prototypes = torch.stack(proto_list)
643 return prototypes
645 @icontract.require(lambda self: len(self.support) > 0)
646 def get_support(self) -> OrderedDict[int, torch.Tensor]:
647 """
648 Method for returning support examples used in training.
650 Returns
651 -------
652 OrderedDict[int, torch.Tensor]
653 Dictionary containing support examples for each class.
654 """
655 return self.support
657 @icontract.require(lambda self: len(self.prototypes) > 0)
658 def get_prototypes(self) -> torch.Tensor:
659 """
660 Method for returning class prototypes.
662 Returns
663 -------
664 torch.Tensor
665 Tensors of prototypes for each of the given classes in the support.
666 """
667 return self.prototypes
669 @icontract.require(lambda num_calibration_epochs: 0 < num_calibration_epochs)
670 @icontract.require(lambda calibration_lr: calibration_lr > 0.0)
671 def calibrate_model(
672 self,
673 dataset: torch.utils.data.Dataset,
674 num_calibration_epochs: int = 1,
675 calibration_lr: float = 0.01,
676 calibration_batch_size: int = 256,
677 ) -> None:
678 """
679 Fine-tune the temperature after training. Note this function is also run at the conclusion of train_model.
681 Parameters
682 ----------
683 dataset : TensorDataset
684 An iterable, pytorch TensorDataset.
685 num_calibration_epochs : int, optional
686 Number of epochs to tune temperature.
687 calibration_lr : float, optional
688 Learning rate for temperature optimization.
689 """
691 calibration_loader = DataLoader(
692 dataset,
693 batch_size=calibration_batch_size,
694 shuffle=True,
695 drop_last=False,
696 )
698 self.temperature.requires_grad = True
699 loss_fn = torch.nn.functional.cross_entropy
700 optimizer = torch.optim.Adam([self.temperature], lr=calibration_lr)
701 for _ in range(num_calibration_epochs):
702 for xs, labels in calibration_loader:
703 optimizer.zero_grad()
704 xs = xs.to(self.device)
705 labels = labels.to(self.device)
706 with torch.no_grad():
707 logits = self.model(xs)
708 logits = logits / self.temperature
709 loss = loss_fn(logits, labels.to(torch.long))
710 loss.backward()
711 optimizer.step()
712 self.temperature.requires_grad = False
714 def forward(self, X: torch.Tensor) -> torch.Tensor:
715 """
716 EquineGP forward function, generates logits for classification.
718 Parameters
719 ----------
720 X : torch.Tensor
721 Input tensor for generating predictions.
723 Returns
724 -------
725 torch.Tensor
726 Output probabilities computed.
727 """
728 X = X.to(self.device)
729 preds = self.model(X)
730 return preds / self.temperature.to(self.device)
732 @icontract.ensure(
733 lambda result: all((0 <= result.ood_scores) & (result.ood_scores <= 1.0))
734 )
735 def predict(self, X: torch.Tensor) -> EquineOutput:
736 """
737 Predict function for EquineGP, inherited and implemented from Equine.
739 Parameters
740 ----------
741 X : torch.Tensor
742 Input tensor.
744 Returns
745 -------
746 EquineOutput
747 Output object containing prediction probabilities and OOD scores.
748 """
749 logits = self(X)
750 preds = torch.softmax(logits, dim=1)
751 equiprobable = torch.ones(self.num_outputs) / self.num_outputs
752 max_entropy = torch.sum(torch.special.entr(equiprobable))
753 ood_score = torch.sum(torch.special.entr(preds), dim=1) / max_entropy
754 embeddings = self.compute_embeddings(X)
755 eq_out = EquineOutput(
756 classes=preds, ood_scores=ood_score, embeddings=embeddings
757 ) # TODO return embeddings
759 self.validate_feature_label_names(X.shape[-1], self.num_outputs)
761 return eq_out
763 def save(self, path: str) -> None:
764 """
765 Function to save all model parameters to a file.
767 Parameters
768 ----------
769 path : str
770 Filename to write the model.
771 """
772 model_settings = {
773 "emb_out_dim": self.num_deep_features,
774 "num_classes": self.num_outputs,
775 "num_random_features": self.num_random_features,
776 "init_temperature": self.temperature.item(),
777 "device": self.device_type,
778 }
780 jit_model = torch.jit.script(prepare_jit_module(self.model.feature_extractor))
781 buffer = io.BytesIO()
782 torch.jit.save(jit_model, buffer)
784 laplace_sd = self.model.state_dict()
785 keys_to_delete = []
786 for key in laplace_sd:
787 if "feature_extractor" in key:
788 keys_to_delete.append(key)
789 for key in keys_to_delete:
790 del laplace_sd[key]
792 # Plain values only, so load_checkpoint can read it with weights_only=True.
793 save_data = {
794 "equine_format_version": EQUINE_FORMAT_VERSION,
795 "embed_jit_save": jit_archive_to_tensor(buffer),
796 "feature_names": self.feature_names,
797 "label_names": self.label_names,
798 "laplace_model_save": laplace_sd,
799 "num_data": self.model.num_data,
800 "settings": model_settings,
801 "support": {int(label): x for label, x in self.support.items()},
802 "train_batch_size": self.model.train_batch_size,
803 "train_summary": self.train_summary,
804 }
806 torch.save(save_data, path) # TODO allow model checkpointing
808 @classmethod
809 def load(cls, path: str) -> Equine:
810 """
811 Load a previously saved EquineGP model.
813 Files from earlier EQUINE releases must first be rewritten with
814 ``python -m equine.upgrade``.
816 Parameters
817 ----------
818 path : str
819 Input filename.
821 Returns
822 -------
823 EquineGP
824 The reconstituted EquineGP object.
826 Raises
827 ------
828 ValueError
829 If the file is not in the current EQUINE format.
830 """
831 return cls._from_checkpoint(load_checkpoint(path))
833 @classmethod
834 def _from_checkpoint(cls, model_save: dict[str, Any]) -> Equine:
835 """
836 Rebuild an EquineGP from an already-loaded checkpoint dictionary.
838 Parameters
839 ----------
840 model_save : dict[str, Any]
841 The dictionary returned by ``utils.load_checkpoint``.
843 Returns
844 -------
845 EquineGP
846 The reconstituted EquineGP object.
847 """
848 jit_model = load_jit_archive(model_save.get("embed_jit_save"))
849 eq_model = cls(jit_model, **model_save.get("settings"))
851 eq_model.feature_names = model_save.get("feature_names")
852 eq_model.label_names = model_save.get("label_names")
853 eq_model.train_summary = model_save.get("train_summary")
855 eq_model.model.load_state_dict(
856 model_save.get("laplace_model_save"), strict=False
857 )
858 eq_model.model.seen_data = model_save.get("laplace_model_save").get("seen_data")
860 eq_model.model.set_training_params(
861 model_save.get("num_data"), model_save.get("train_batch_size")
862 )
863 eq_model.eval()
865 support = OrderedDict(
866 (int(label), x) for label, x in model_save.get("support").items()
867 )
868 if len(support) > 0:
869 eq_model.support = support
870 eq_model.prototypes = eq_model.compute_prototypes()
872 return eq_model