Coverage for /opt/hostedtoolcache/Python/3.12.15/x64/lib/python3.12/site-packages/equine/utils.py: 100%
196 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 pickle
7import sys
8from collections import OrderedDict
9from typing import Any, Optional, Union
11import icontract
12import torch
13from beartype import beartype
14from torchmetrics.classification import (
15 MulticlassAccuracy,
16 MulticlassCalibrationError,
17 MulticlassConfusionMatrix,
18 MulticlassF1Score,
19)
21from .equine import Equine
22from .equine_output import EquineOutput
24# Layout version written by ``save``. Version 2 holds only tensors, containers
25# and plain scalars/strings, so ``torch.load(weights_only=True)`` reads it.
26EQUINE_FORMAT_VERSION = 2
28# torch < 2.6 has a known bypass of weights_only=True (CVE-2025-32434).
29_MIN_SAFE_TORCH = (2, 6)
31_LEGACY_FILE_ERROR = (
32 "'{path}' is not in the current EQUINE model format: it has no "
33 "equine_format_version, or it holds pickled Python objects, as files saved "
34 "by earlier EQUINE releases do. Rewrite it once with "
35 "`python -m equine.upgrade {path}`. Only upgrade files you trust: the upgrade "
36 "unpickles the file without restrictions, which can run code embedded in it."
37)
40class LegacyFormatError(ValueError):
41 """The file predates the current format; ``python -m equine.upgrade`` rewrites it."""
44def _torch_version() -> tuple[int, int]:
45 major, minor = (int(p) for p in torch.__version__.split("+")[0].split(".")[:2])
46 return (major, minor)
49def load_checkpoint(path: str, map_location: Optional[str] = None) -> dict[str, Any]:
50 """
51 Read a saved EQUINE model file with ``torch.load(weights_only=True)``.
53 Files from earlier releases hold pickled Python objects and are refused;
54 ``python -m equine.upgrade`` rewrites them in the current format.
56 Parameters
57 ----------
58 path : str
59 Filename of the saved model.
60 map_location : Optional[str]
61 Device to map the saved tensors onto (passed to ``torch.load``).
63 Returns
64 -------
65 dict[str, Any]
66 The saved checkpoint dictionary.
68 Raises
69 ------
70 ValueError
71 If torch is older than 2.6, if the file cannot be read as a PyTorch
72 file, if it is not in the current format (a ``LegacyFormatError``), or
73 if it declares a format version this EQUINE does not know.
74 """
75 if _torch_version() < _MIN_SAFE_TORCH:
76 raise ValueError(
77 "EQUINE needs torch >= 2.6 to read model files (installed: "
78 f"{torch.__version__}): torch.load(weights_only=True) has a known "
79 "bypass on earlier releases (CVE-2025-32434)."
80 )
81 try:
82 checkpoint = torch.load(path, map_location=map_location, weights_only=True)
83 except pickle.UnpicklingError as err:
84 raise LegacyFormatError(_LEGACY_FILE_ERROR.format(path=path)) from err
85 except Exception as err: # not a PyTorch file at all
86 raise ValueError(
87 f"'{path}' could not be read as a PyTorch file "
88 f"({type(err).__name__}: {err})."
89 ) from err
90 if not isinstance(checkpoint, dict) or "equine_format_version" not in checkpoint:
91 raise LegacyFormatError(_LEGACY_FILE_ERROR.format(path=path))
92 version = checkpoint["equine_format_version"]
93 if version > EQUINE_FORMAT_VERSION:
94 raise ValueError(
95 f"Unsupported EQUINE model format version {version}; this EQUINE "
96 f"supports up to {EQUINE_FORMAT_VERSION}. Upgrade EQUINE."
97 )
98 return checkpoint
101def jit_archive_to_tensor(buffer: io.BytesIO) -> torch.Tensor:
102 """A serialized TorchScript archive as a ``uint8`` tensor, for ``save``."""
103 return torch.frombuffer(bytearray(buffer.getvalue()), dtype=torch.uint8)
106def _archive_bytes(archive: torch.Tensor) -> bytes:
107 """The bytes of a ``uint8`` archive tensor (exactly the tensor's own view)."""
108 return archive.detach().cpu().contiguous().numpy().tobytes()
111def load_jit_archive(
112 archive: torch.Tensor, map_location: Optional[str] = None
113) -> torch.jit.ScriptModule:
114 """Reconstruct the TorchScript module stored by ``jit_archive_to_tensor``."""
115 buffer = io.BytesIO(_archive_bytes(archive))
116 return torch.jit.load(buffer, map_location=map_location)
119@icontract.ensure(lambda result, module: result is module)
120@beartype
121def prepare_jit_module(module: torch.nn.Module) -> torch.nn.Module:
122 """
123 Make an ``nn.Module`` safe to pass to ``torch.jit.script`` on Python 3.14+.
125 Starting with Python 3.14 (PEP 649/749) a class's ``__annotations__`` is
126 exposed through a ``type`` getset descriptor instead of being stored in the
127 class ``__dict__``. ``torch.jit``'s scripting checker reads
128 ``__annotations__`` from the module *instance*, where attribute lookup only
129 consults the ``__dict__`` of each class in the MRO. For any ``nn.Module``
130 that declares no class-level annotations this lookup misses and
131 ``nn.Module.__getattr__`` raises ``AttributeError``, breaking
132 ``torch.jit.script``. Materializing each submodule's own class annotations
133 onto its instance ``__dict__`` lets that lookup succeed. This is a no-op on
134 Python < 3.14.
136 Parameters
137 ----------
138 module : torch.nn.Module
139 The module that is about to be scripted with ``torch.jit.script``.
141 Returns
142 -------
143 torch.nn.Module
144 The same module, returned for convenient inline use.
145 """
146 if sys.version_info >= (3, 14):
147 for submodule in module.modules():
148 if "__annotations__" not in submodule.__dict__:
149 submodule.__dict__["__annotations__"] = dict(
150 type(submodule).__annotations__
151 )
152 return module
155@icontract.require(lambda y_hat, y_test: y_hat.size(dim=0) == y_test.size(dim=0))
156@icontract.ensure(lambda result: result >= 0.0)
157@beartype
158def brier_score(y_hat: torch.Tensor, y_test: torch.Tensor) -> float:
159 """
160 Compute the Brier score for a multiclass problem:
161 $$ \\frac{1}{N} \\sum_{i=1}^{N} \\sum_{j=1}^{M} (f_{ij} - o_{ij})^2 , $$
162 where $f_{ij}$ is the predicted probability of class $j$ for inference sample $i$
163 and $o_{ij}$ is the one-hot encoded ground truth label.
165 Parameters
166 ----------
167 y_hat : torch.Tensor
168 Probabilities for each class.
169 y_test : torch.Tensor
170 Integer argument class labels (ground truth).
172 Returns
173 -------
174 float
175 Brier score.
176 """
177 _, num_classes = y_hat.size()
178 one_hot_y_test = torch.nn.functional.one_hot(y_test.long(), num_classes=num_classes)
179 bs = torch.mean(torch.sum((y_hat - one_hot_y_test) ** 2, dim=1)).item()
180 return bs
183@icontract.require(lambda y_hat, y_test: y_hat.size(dim=0) == y_test.size(dim=0))
184@icontract.ensure(lambda result: result <= 1.0)
185@beartype
186def brier_skill_score(y_hat: torch.Tensor, y_test: torch.Tensor) -> float:
187 """
188 Compute the Brier skill score as compared to randomly guessing.
190 Parameters
191 ----------
192 y_hat : torch.Tensor
193 Probabilities for each class.
194 y_test : torch.Tensor
195 Integer argument class labels (ground truth).
197 Returns
198 -------
199 float
200 Brier skill score.
201 """
202 _, num_classes = y_hat.size()
203 random_guess = (1.0 / num_classes) * torch.ones(y_hat.size())
204 bs0 = brier_score(random_guess, y_test)
205 bs1 = brier_score(y_hat, y_test)
206 bss = 1.0 - bs1 / bs0
207 return bss
210@icontract.require(lambda y_hat, y_test: y_hat.size(dim=0) == y_test.size(dim=0))
211@icontract.ensure(lambda result: (0.0 <= result) and (result <= 1.0))
212@beartype
213def expected_calibration_error(y_hat: torch.Tensor, y_test: torch.Tensor) -> float:
214 """
215 Compute the expected calibration error (ECE) for a multiclass problem.
217 Parameters
218 ----------
219 y_hat : torch.Tensor
220 Probabilities for each class.
221 y_test : torch.Tensor
222 Class label indices (ground truth).
224 Returns
225 -------
226 float
227 Expected calibration error.
228 """
229 _, num_classes = y_hat.size()
230 metric = MulticlassCalibrationError(num_classes=num_classes, n_bins=25, norm="l1")
231 ece = metric(y_hat, y_test).item()
232 return ece
235@icontract.require(
236 lambda train_y, selected_labels: len(selected_labels) <= len(train_y)
237)
238@icontract.ensure(
239 lambda result, selected_labels: set(result.keys()).issubset(set(selected_labels))
240)
241@beartype
242def _get_shuffle_idxs_by_class(
243 train_y: torch.Tensor, selected_labels: list
244) -> dict[Any, torch.Tensor]:
245 """
246 Internal helper function to randomly select indices of example classes for a given
247 set of labels.
249 Parameters
250 ----------
251 train_y : torch.Tensor
252 Label data.
253 selected_labels : list
254 list of unique labels found in the label data.
256 Returns
257 -------
258 dict[Any, torch.Tensor]
259 Tensor of indices corresponding to each label.
260 """
261 shuffled_idxs_by_class = OrderedDict()
262 for label in selected_labels:
263 label_idxs = torch.argwhere(train_y == label).squeeze()
264 shuffled_idxs_by_class[label] = label_idxs[torch.randperm(label_idxs.shape[0])]
266 return shuffled_idxs_by_class
269@icontract.require(lambda train_x, train_y: len(train_x) <= len(train_y))
270@icontract.require(
271 lambda selected_labels, train_x: (
272 (0 < len(selected_labels)) & (len(selected_labels) < len(train_x))
273 )
274)
275@icontract.require(
276 lambda support_size, train_x: (0 < support_size) & (support_size < len(train_x))
277)
278@icontract.require(
279 lambda support_size, selected_labels, train_x: (
280 support_size * len(selected_labels) <= len(train_x)
281 )
282)
283@icontract.require(
284 lambda selected_labels, shuffled_indexes: (
285 (len(shuffled_indexes.keys()) == len(selected_labels))
286 if shuffled_indexes is not None
287 else True
288 )
289)
290@icontract.ensure(
291 lambda result, selected_labels: len(result.keys()) == len(selected_labels)
292)
293@beartype
294def generate_support(
295 train_x: torch.Tensor,
296 train_y: torch.Tensor,
297 support_size: int,
298 selected_labels: list[Any],
299 shuffled_indexes: Union[None, dict[Any, torch.Tensor]] = None,
300) -> OrderedDict[int, torch.Tensor]:
301 """
302 Randomly select `support_size` examples of `way` classes from the examples in
303 `train_x` with corresponding labels in `train_y` and return them as a dictionary.
305 Parameters
306 ----------
307 train_x : torch.Tensor
308 Input training data.
309 train_y : torch.Tensor
310 Corresponding classification labels.
311 support_size : int
312 Number of support examples for each class.
313 selected_labels : list
314 Selected class labels to generate examples from.
315 shuffled_indexes: Union[None, dict[Any, torch.Tensor]], optional
316 Simply use the precomputed indexes if they are available
318 Returns
319 -------
320 OrderedDict[int, torch.Tensor]
321 Ordered dictionary of class labels with corresponding support examples.
322 """
323 labels, counts = torch.unique(train_y, return_counts=True)
324 if shuffled_indexes is None:
325 for label, count in list(zip(labels, counts)):
326 if (label in selected_labels) and (count < support_size):
327 raise ValueError(f"Not enough support examples in class {label}")
328 shuffled_idxs = _get_shuffle_idxs_by_class(train_y, selected_labels)
329 else:
330 shuffled_idxs = shuffled_indexes
332 support = OrderedDict[int, torch.Tensor]()
333 for label in selected_labels:
334 shuffled_x = train_x[shuffled_idxs[label]]
336 assert torch.unique(train_y[shuffled_idxs[label]]).tolist() == [label], (
337 "Not enough support for label " + str(label)
338 )
339 selected_support = shuffled_x[:support_size]
340 support[int(label)] = selected_support
342 return support
345@icontract.require(lambda train_x: len(train_x.shape) >= 2)
346@icontract.require(lambda train_y: len(train_y.shape) == 1)
347@icontract.require(lambda support_size: support_size > 1)
348@icontract.require(lambda way: way > 0)
349@icontract.require(lambda episode_size: episode_size > 0)
350@icontract.ensure(lambda result: len(result) == 3)
351@icontract.ensure(lambda result: result[1].shape[0] == result[2].shape[0])
352@icontract.ensure(lambda way, result: len(result[0]) == way)
353@icontract.ensure(
354 lambda support_size, result: all(
355 len(support) == support_size for support in result[0].values()
356 )
357)
358@beartype
359def generate_episode(
360 train_x: torch.Tensor,
361 train_y: torch.Tensor,
362 support_size: int,
363 way: int,
364 episode_size: int,
365) -> tuple[OrderedDict[int, torch.Tensor], torch.Tensor, torch.Tensor]:
366 """
367 Generate a single episode of data for a few-shot learning task.
369 Parameters
370 ----------
371 train_x : torch.Tensor
372 Input training data.
373 train_y : torch.Tensor
374 Corresponding classification labels.
375 support_size : int
376 Number of support examples for each class.
377 way : int
378 Number of classes in the episode.
379 episode_size : int
380 Total number of examples in the episode.
382 Returns
383 -------
384 tuple[dict[Any, torch.Tensor], torch.Tensor, torch.Tensor]
385 tuple of support examples, query examples, and query labels.
386 """
387 labels, counts = torch.unique(train_y, return_counts=True)
388 if way > len(labels):
389 raise ValueError(
390 f"The way (#classes in each episode), {way}, must be <= number of labels, {len(labels)}"
391 )
393 selected_labels = sorted(
394 labels[torch.randperm(labels.shape[0])][:way].tolist()
395 ) # need to be in same order every time
397 for label, count in list(zip(labels, counts)):
398 if (label in selected_labels) and (count < support_size):
399 raise ValueError(f"Not enough support examples in class {label}")
400 shuffled_idxs = _get_shuffle_idxs_by_class(train_y, selected_labels)
402 support = generate_support(
403 train_x, train_y, support_size, selected_labels, shuffled_idxs
404 )
406 examples_per_task = episode_size // way
408 episode_data_list = []
409 episode_label_list = []
410 episode_support = OrderedDict()
411 for episode_label, label in enumerate(selected_labels):
412 shuffled_x = train_x[shuffled_idxs[label]]
413 shuffled_y = torch.Tensor(
414 [episode_label] * len(shuffled_idxs[label])
415 ) # need sequential labels for episode
417 num_remaining_examples = shuffled_x.shape[0] - support_size
418 assert num_remaining_examples > 0, (
419 "Cannot have "
420 + str(num_remaining_examples)
421 + " left with support_size "
422 + str(support_size)
423 + " and shape "
424 + str(shuffled_x.shape)
425 + " from train_x shaped "
426 + str(train_x.shape)
427 )
428 episode_end_idx = support_size + min(num_remaining_examples, examples_per_task)
430 episode_data_list.append(shuffled_x[support_size:episode_end_idx])
431 episode_label_list.append(shuffled_y[support_size:episode_end_idx])
432 episode_support[episode_label] = support[label]
434 episode_x = torch.concat(episode_data_list)
435 episode_y = torch.concat(episode_label_list)
437 return episode_support, episode_x, episode_y.squeeze().to(torch.long)
440@icontract.require(
441 lambda eq_preds, true_y: eq_preds.classes.size(dim=0) == true_y.size(dim=0)
442)
443@beartype
444def generate_model_metrics(
445 eq_preds: EquineOutput, true_y: torch.Tensor
446) -> dict[str, Any]:
447 """
448 Generate various metrics for evaluating a model's performance.
450 Parameters
451 ----------
452 eq_preds : EquineOutput
453 Model predictions.
454 true_y : torch.Tensor
455 True class labels.
457 Returns
458 -------
459 dict[str, Any]
460 Dictionary of model metrics.
461 """
462 pred_y = torch.argmax(eq_preds.classes, dim=1)
463 accuracy = MulticlassAccuracy(num_classes=eq_preds.classes.shape[1])
464 f1_score = MulticlassF1Score(num_classes=eq_preds.classes.shape[1], average="micro")
465 confusion_matrix = MulticlassConfusionMatrix(num_classes=eq_preds.classes.shape[1])
466 metrics = {
467 "accuracy": accuracy(true_y, pred_y),
468 "microF1Score": f1_score(true_y, pred_y),
469 "confusionMatrix": confusion_matrix(true_y, pred_y).tolist(),
470 "brierScore": brier_score(eq_preds.classes, true_y),
471 "brierSkillScore": brier_skill_score(eq_preds.classes, true_y),
472 "expectedCalibrationError": expected_calibration_error(
473 eq_preds.classes, true_y
474 ),
475 }
476 return metrics
479@icontract.require(lambda Y: len(Y.shape) == 1)
480@icontract.ensure(
481 lambda result: all("label" in d and "numExamples" in d for d in result)
482)
483@icontract.ensure(lambda result: all(d["numExamples"] >= 0 for d in result))
484@beartype
485def get_num_examples_per_label(Y: torch.Tensor) -> list[dict[str, Any]]:
486 """
487 Get the number of examples per label in the given tensor.
489 Parameters
490 ----------
491 Y : torch.Tensor
492 Tensor of class labels.
494 Returns
495 -------
496 list[dict[str, Any]]
497 list of dictionaries containing label and number of examples.
498 """
499 tensor_labels, tensor_counts = Y.unique(return_counts=True)
501 examples_per_label = []
502 for i, label in enumerate(tensor_labels):
503 examples_per_label.append(
504 {"label": label.item(), "numExamples": tensor_counts[i].item()}
505 )
507 return examples_per_label
510@icontract.require(lambda train_y: train_y.shape[0] > 0)
511@beartype
512def generate_train_summary(
513 model: Equine, train_y: torch.Tensor, date_trained: str
514) -> dict[str, Any]:
515 """
516 Generate a summary of the training data.
518 Parameters
519 ----------
520 model : Equine
521 Model object.
522 train_y : torch.Tensor
523 Training labels.
524 date_trained : str
525 Date of training.
527 Returns
528 -------
529 dict[str, Any]
530 Dictionary containing training summary.
531 """
532 train_summary = {
533 "numTrainExamples": get_num_examples_per_label(train_y),
534 "dateTrained": date_trained,
535 "modelType": model.__class__.__name__,
536 }
537 return train_summary
540@icontract.require(
541 lambda eq_preds, test_y: test_y.shape[0] == eq_preds.classes.shape[0]
542)
543@beartype
544def generate_model_summary(
545 model: Equine,
546 eq_preds: EquineOutput,
547 test_y: torch.Tensor,
548) -> dict[str, Any]:
549 """
550 Generate a summary of the model's performance.
552 Parameters
553 ----------
554 model : Equine
555 Model object.
556 eq_preds : EquineOutput
557 Model predictions.
558 test_y : torch.Tensor
559 True class labels.
561 Returns
562 -------
563 dict[str, Any]
564 Dictionary containing model summary.
565 """
566 summary = generate_model_metrics(eq_preds, test_y)
567 summary["numTestExamples"] = get_num_examples_per_label(test_y)
568 summary.update(model.train_summary) # union of train_summary and generated metrics
570 return summary
573@icontract.require(lambda cov: cov.shape[-2] == cov.shape[-1])
574def mahalanobis_distance_nosq(x: torch.Tensor, cov: torch.Tensor) -> torch.Tensor:
575 """
576 Compute Mahalanobis distance $x^T C x$ (without square root), assume cov is symmetric positive definite
578 Parameters
579 ----------
580 x : torch.Tensor
581 vectors to compute distances for
582 cov : torch.Tensor
583 covariance matrix, assumes first dimension is number of classes
584 """
585 U, S, _ = torch.linalg.svd(cov)
586 S_inv_sqrt = torch.stack(
587 [torch.diag(torch.sqrt(1.0 / S[i])) for i in range(S.shape[0])], dim=0
588 )
589 prod = torch.matmul(S_inv_sqrt, torch.transpose(U, 1, 2))
590 dist = torch.sum(torch.square(torch.matmul(prod, x)), dim=1)
591 return dist
594@icontract.require(
595 lambda X, Y: X.shape[0] == Y.shape[0],
596 "X and Y must have the same number of samples.",
597)
598@icontract.require(
599 lambda test_size: 0.0 < test_size < 1.0, "test_size must be between 0 and 1."
600)
601@icontract.ensure(
602 lambda result: len(result) == 4, "Function must return four elements."
603)
604@icontract.ensure(
605 lambda X, result: result[0].shape[0] + result[1].shape[0] == X.shape[0],
606 "Total samples must be preserved.",
607)
608@icontract.ensure(
609 lambda Y, result: result[2].shape[0] + result[3].shape[0] == Y.shape[0],
610 "Total labels must be preserved.",
611)
612@icontract.ensure(
613 lambda result: result[0].shape[0] == result[2].shape[0],
614 "Train features and labels must match in size.",
615)
616@icontract.ensure(
617 lambda result: result[1].shape[0] == result[3].shape[0],
618 "Test features and labels must match in size.",
619)
620@beartype
621def stratified_train_test_split(
622 X: torch.Tensor, Y: torch.Tensor, test_size: float
623) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
624 """
625 A pytorch-ified version of sklearn's train_test_split with data stratification
627 Parameters
628 ----------
629 X : torch.Tensor
630 Input features tensor of shape (n_samples, n_features).
631 Y : torch.Tensor
632 Labels tensor of shape (n_samples,).
633 test_size : float
634 Proportion of the dataset to include in the test split (between 0.0 and 1.0).
636 Returns
637 -------
638 train_x : torch.Tensor
639 Training set features.
640 calib_x : torch.Tensor
641 Test set features.
642 train_y : torch.Tensor
643 Training set labels.
644 calib_y : torch.Tensor
645 Test set labels.
646 """
647 unique_classes, class_counts = torch.unique(Y, return_counts=True)
648 test_counts = (class_counts.float() * test_size).round().long()
649 train_indices = []
650 test_indices = []
652 for cls, test_count in zip(unique_classes, test_counts):
653 cls_indices = torch.where(Y == cls)[0]
654 cls_indices = cls_indices[torch.randperm(len(cls_indices))]
655 test_idx = cls_indices[:test_count]
656 train_idx = cls_indices[test_count:]
657 train_indices.append(train_idx)
658 test_indices.append(test_idx)
660 train_indices = torch.cat(train_indices)
661 test_indices = torch.cat(test_indices)
663 train_x = X[train_indices]
664 train_y = Y[train_indices]
665 calib_x = X[test_indices]
666 calib_y = Y[test_indices]
668 return train_x, calib_x, train_y, calib_y