Coverage for /opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/equine/equine_protonet.py: 96%

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 

5 

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 

13 

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 

21 

22from .equine import Equine, EquineOutput 

23from .utils import ( 

24 generate_episode, 

25 generate_support, 

26 generate_train_summary, 

27 mahalanobis_distance_nosq, 

28 prepare_jit_module, 

29 stratified_train_test_split, 

30) 

31 

32 

33##################################### 

34class CovType(Enum): 

35 """ 

36 Enum class for covariance types used in EQUINE. 

37 """ 

38 

39 UNIT = "unit" 

40 DIAGONAL = "diag" 

41 FULL = "full" 

42 

43 

44PRED_COV_TYPE = CovType.DIAGONAL 

45OOD_COV_TYPE = CovType.DIAGONAL 

46DEFAULT_EPSILON = 1e-5 

47COV_REG_TYPE = "epsilon" 

48 

49 

50############################################### 

51 

52 

53@beartype 

54class Protonet(torch.nn.Module): 

55 """ 

56 Private class that implements a prototypical neural network for use in EQUINE. 

57 """ 

58 

59 def __init__( 

60 self, 

61 embedding_model: torch.nn.Module, 

62 emb_out_dim: int, 

63 cov_type: CovType, 

64 cov_reg_type: str, 

65 epsilon: float, 

66 device: str = "cpu", 

67 ) -> None: 

68 """ 

69 Protonet class constructor. 

70 

71 Parameters 

72 ---------- 

73 embedding_model : torch.nn.Module 

74 The PyTorch embedding model to generate logits with. 

75 emb_out_dim : int 

76 Dimension size of given embedding model's output. 

77 cov_type : CovType 

78 Type of covariance to use when computing distances [unit, diag, full]. 

79 cov_reg_type : str 

80 Type of regularization to use when generating the covariance matrix [epsilon, shared]. 

81 epsilon : float 

82 Epsilon value to use for covariance regularization. 

83 device : str, optional 

84 The device to train the protonet model on (defaults to cpu). 

85 """ 

86 super().__init__() 

87 self.embedding_model = embedding_model 

88 self.cov_type = cov_type 

89 self.cov_reg_type = cov_reg_type 

90 self.epsilon = epsilon 

91 self.emb_out_dim = emb_out_dim 

92 self.to(device) 

93 self.device = device 

94 

95 self.support: OrderedDict[int, torch.Tensor] = OrderedDict() 

96 self.support_embeddings: OrderedDict[int, torch.Tensor] = OrderedDict() 

97 self.model_head: torch.nn.Module = self.create_model_head(emb_out_dim) 

98 self.model_head.to(device) 

99 

100 def create_model_head(self, emb_out_dim: int) -> torch.nn.Identity: 

101 """ 

102 Method for adding a PyTorch layer on top of the given embedding model. This layer 

103 is intended to offer extra degrees of freedom for distance learning in the embedding space. 

104 

105 Parameters 

106 ---------- 

107 emb_out_dim : int 

108 Dimension size of the embedding model output. 

109 

110 Returns 

111 ------- 

112 torch.nn.Linear 

113 The created PyTorch model layer. 

114 """ 

115 return torch.nn.Identity() # (emb_out_dim, emb_out_dim) 

116 

117 def compute_embeddings(self, X: torch.Tensor) -> torch.Tensor: 

118 """ 

119 Method for calculating model embeddings using both the given embedding model and the added model head. 

120 

121 Parameters 

122 ---------- 

123 X : torch.Tensor 

124 Input tensor to compute embeddings on. 

125 

126 Returns 

127 ------- 

128 torch.Tensor 

129 Fully computed embedding tensors for the given X tensor. 

130 """ 

131 model_embeddings = self.embedding_model(X.to(self.device)) 

132 head_embeddings = self.model_head(model_embeddings) 

133 return head_embeddings 

134 

135 @icontract.require(lambda self: len(self.support_embeddings) > 0) 

136 def compute_prototypes(self) -> torch.Tensor: 

137 """ 

138 Method for computing class prototypes based on given support examples. 

139 ``Prototypes'' in this context are the means of the support embeddings for each class. 

140 

141 Returns 

142 ------- 

143 torch.Tensor 

144 Tensors of prototypes for each of the given classes in the support. 

145 """ 

146 # Compute prototype for each class 

147 proto_list = [] 

148 for label in self.support_embeddings: # look at doing functorch 

149 class_prototype = torch.mean(self.support_embeddings[label], dim=0) 

150 proto_list.append(class_prototype) 

151 

152 prototypes = torch.stack(proto_list) 

153 

154 return prototypes 

155 

156 @icontract.require(lambda self: len(self.support_embeddings) > 0) 

157 def compute_covariance(self, cov_type: CovType) -> torch.Tensor: 

158 """ 

159 Method for generating the (regularized) support example covariance matrix(es) used for calculating distances. 

160 Note that this method is only called once per episode, and the resulting tensor is used for all queries. 

161 

162 Parameters 

163 ---------- 

164 cov_type : CovType 

165 Type of covariance to use [unit, diag, full]. 

166 

167 Returns 

168 ------- 

169 torch.Tensor 

170 Tensor containing the generated regularized covariance matrix. 

171 """ 

172 class_cov_dict = OrderedDict().fromkeys( 

173 self.support_embeddings.keys(), torch.Tensor() 

174 ) 

175 for label in self.support_embeddings.keys(): 

176 class_covariance = self.compute_covariance_by_type( 

177 cov_type, self.support_embeddings[label] 

178 ) 

179 class_cov_dict[label] = class_covariance 

180 

181 reg_covariance_dict = self.regularize_covariance( 

182 class_cov_dict, cov_type, self.cov_reg_type 

183 ) 

184 reg_covariance = torch.stack(list(reg_covariance_dict.values())) 

185 

186 return reg_covariance # TODO try putting everything on GPU with .to() and see if faster 

187 

188 def compute_covariance_by_type( 

189 self, cov_type: CovType, embedding: torch.Tensor 

190 ) -> torch.Tensor: 

191 """ 

192 Select the appropriate covariance matrix type based on cov_type. 

193 

194 Parameters 

195 ---------- 

196 cov_type : str 

197 Type of covariance to use. Options are ['unit', 'diag', 'full']. 

198 embedding : torch.Tensor 

199 Embedding tensor to use when generating the covariance matrix. 

200 

201 Returns 

202 ------- 

203 torch.Tensor 

204 Tensor containing the requested covariance matrix. 

205 """ 

206 if cov_type == CovType.FULL: 

207 class_covariance = torch.cov(embedding.T) 

208 elif cov_type == CovType.DIAGONAL: 

209 class_covariance = torch.var(embedding, dim=0) 

210 elif cov_type == CovType.UNIT: 

211 class_covariance = torch.ones(self.emb_out_dim) 

212 else: 

213 raise ValueError 

214 

215 return class_covariance 

216 

217 def regularize_covariance( 

218 self, 

219 class_cov_dict: OrderedDict[int, torch.Tensor], 

220 cov_type: CovType, 

221 cov_reg_type: str, 

222 ) -> OrderedDict[int, torch.Tensor]: 

223 """ 

224 Method to add regularization to each class covariance matrix based on the selected regularization type. 

225 

226 Parameters 

227 ---------- 

228 class_cov_dict : OrderedDict[int, torch.Tensor] 

229 A dictionary containing each class and the corresponding covariance matrix. 

230 cov_type : CovType 

231 Type of covariance to use [unit, diag, full]. 

232 

233 Returns 

234 ------- 

235 dict[float, torch.Tensor] 

236 Dictionary containing the regularized class covariance matrices. 

237 """ 

238 

239 if cov_type == CovType.FULL: 

240 regularization = torch.diag(self.epsilon * torch.ones(self.emb_out_dim)).to( 

241 self.device 

242 ) 

243 elif cov_type == CovType.DIAGONAL: 

244 regularization = self.epsilon * torch.ones(self.emb_out_dim).to(self.device) 

245 elif cov_type == CovType.UNIT: 245 ↛ 248line 245 didn't jump to line 248 because the condition on line 245 was always true

246 regularization = torch.zeros(self.emb_out_dim).to(self.device) 

247 

248 if cov_reg_type == "shared": 

249 if cov_type != CovType.FULL and cov_type != CovType.DIAGONAL: 249 ↛ 250line 249 didn't jump to line 250 because the condition on line 249 was never true

250 for label in self.support_embeddings: 

251 class_cov_dict[label] = class_cov_dict[label] + regularization 

252 warnings.warn( 

253 "Covariance type UNIT is incompatible with shared regularization, \ 

254 reverting to epsilon regularization" 

255 ) 

256 return class_cov_dict 

257 

258 shared_covariance = self.compute_shared_covariance(class_cov_dict, cov_type) 

259 

260 for label in self.support_embeddings: 

261 num_class_support = self.support_embeddings[label].shape[0] 

262 lamb = num_class_support / (num_class_support + 1) 

263 

264 class_cov_dict[label] = ( 

265 lamb * class_cov_dict[label] 

266 + (1 - lamb) * shared_covariance 

267 + regularization 

268 ) 

269 

270 elif cov_reg_type == "epsilon": 270 ↛ 276line 270 didn't jump to line 276 because the condition on line 270 was always true

271 for label in class_cov_dict.keys(): 

272 class_cov_dict[label] = ( 

273 class_cov_dict[label].to(self.device) + regularization 

274 ) 

275 

276 return class_cov_dict 

277 

278 def compute_shared_covariance( 

279 self, class_cov_dict: OrderedDict[int, torch.Tensor], cov_type: CovType 

280 ) -> torch.Tensor: 

281 """ 

282 Method to calculate a shared covariance matrix. 

283 

284 The shared covariance matrix is calculated as the weighted average of the class covariance matrices, 

285 where the weights are the number of support examples for each class. This is useful when the number of 

286 support examples for each class is small. 

287 

288 Parameters 

289 ---------- 

290 class_cov_dict : OrderedDict[int, torch.Tensor] 

291 A dictionary containing each class and the corresponding covariance matrix. 

292 cov_type : CovType 

293 Type of covariance to use [unit, diag, full]. 

294 

295 Returns 

296 ------- 

297 torch.Tensor 

298 Tensor containing the shared covariance matrix. 

299 """ 

300 total_support = sum([x.shape[0] for x in class_cov_dict.values()]) 

301 

302 if cov_type == CovType.FULL: 302 ↛ 303line 302 didn't jump to line 303 because the condition on line 302 was never true

303 shared_covariance = torch.zeros((self.emb_out_dim, self.emb_out_dim)) 

304 elif cov_type == CovType.DIAGONAL: 

305 shared_covariance = torch.zeros(self.emb_out_dim) 

306 else: 

307 raise ValueError( 

308 "Shared covariance can only be used with FULL or DIAGONAL (not UNIT) covariance types" 

309 ) 

310 

311 for label in class_cov_dict: 

312 num_class_support = class_cov_dict[label].shape[0] 

313 shared_covariance = ( 

314 shared_covariance + (num_class_support - 1) * class_cov_dict[label] 

315 ) # undo N-1 div from cov 

316 

317 shared_covariance = shared_covariance / ( 

318 total_support - 1 

319 ) # redo N-1 div for shared cov 

320 

321 return shared_covariance 

322 

323 @icontract.require(lambda X_embed, mu: X_embed.shape[-1] == mu.shape[-1]) 

324 @icontract.ensure(lambda result: torch.all(result >= 0)) 

325 def compute_distance( 

326 self, X_embed: torch.Tensor, mu: torch.Tensor, cov: torch.Tensor 

327 ) -> torch.Tensor: 

328 """ 

329 Method to compute the distances to class prototypes for the given embeddings. 

330 

331 Parameters 

332 ---------- 

333 X_embed : torch.Tensor 

334 The embeddings of the query examples. 

335 mu : torch.Tensor 

336 The class prototypes (means of the support embeddings). 

337 cov : torch.Tensor 

338 The support covariance matrix. 

339 

340 Returns 

341 ------- 

342 torch.Tensor 

343 The calculated distances from each of the class prototypes for the given embeddings. 

344 """ 

345 _queries = torch.unsqueeze(X_embed, 1) # examples x 1 x dimension 

346 diff = torch.sub(mu, _queries) 

347 

348 if len(cov.shape) == 2: # (diagonal covariance) 

349 # examples x classes x dimension 

350 sq_diff = diff**2 

351 div = torch.div(sq_diff.to(self.device), cov.to(self.device)) 

352 dist = torch.nan_to_num(div) 

353 dist = torch.sum(dist, dim=2) # examples x classes 

354 dist = dist.squeeze(dim=1) 

355 dist = torch.sqrt(dist + self.epsilon) # examples x classes 

356 else: # len(cov.shape) == 3: (full covariance) 

357 diff = diff.permute(1, 2, 0) # classes x dimension x examples 

358 dist = mahalanobis_distance_nosq(diff, cov) 

359 dist = torch.sqrt(dist.permute(1, 0) + self.epsilon) # examples x classes 

360 dist = dist.squeeze(dim=1) 

361 return dist 

362 

363 def compute_classes(self, distances: torch.Tensor) -> torch.Tensor: 

364 """ 

365 Method to compute predicted classes from distances via a softmax function. 

366 

367 Parameters 

368 ---------- 

369 distances : torch.Tensor 

370 The distances of embeddings to class prototypes. 

371 

372 Returns 

373 ------- 

374 torch.Tensor 

375 Tensor of class predictions. 

376 """ 

377 softmax = torch.nn.functional.softmax(torch.neg(distances), dim=-1) 

378 return softmax 

379 

380 def forward(self, X: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: 

381 """ 

382 Protonet forward function, generates class probability predictions and distances from prototypes. 

383 

384 Parameters 

385 ---------- 

386 X : torch.Tensor 

387 Input tensor of queries for generating predictions. 

388 

389 Returns 

390 ------- 

391 tuple[torch.Tensor, torch.Tensor] 

392 tuple containing class probability predictions, and class distances from prototypes. 

393 """ 

394 if len(self.support) == 0 or len(self.support_embeddings) == 0: 

395 raise ValueError( 

396 "No support examples found. Protonet Model requires model support to \ 

397 be set with the 'update_support()' method before calling forward." 

398 ) 

399 

400 X_embed = self.compute_embeddings(X) 

401 if X_embed.shape == torch.Size([self.emb_out_dim]): 

402 X_embed = X_embed.unsqueeze(dim=0) # handle single examples 

403 distances = self.compute_distance(X_embed, self.prototypes, self.covariance) 

404 classes = self.compute_classes(distances) 

405 

406 return classes, distances 

407 

408 def update_support(self, support: OrderedDict[int, torch.Tensor]) -> None: 

409 """ 

410 Method to update the support examples, and all the calculations that rely on them. 

411 

412 Parameters 

413 ---------- 

414 support : OrderedDict 

415 Ordered dict containing class labels and their associated support examples. 

416 """ 

417 self.support = support # TODO torch.nn.ParameterDict(support) 

418 

419 support_embs = OrderedDict().fromkeys(support.keys(), torch.Tensor()) 

420 for label in support: 

421 support_embs[label] = self.compute_embeddings(support[label]) 

422 

423 self.support_embeddings = ( 

424 support_embs # TODO torch.nn.ParameterDict(support_embs) 

425 ) 

426 

427 self.prototypes: torch.Tensor = self.compute_prototypes() 

428 

429 if self.training is False: 

430 self.compute_global_moments() 

431 self.covariance: torch.Tensor = self.compute_covariance( 

432 cov_type=PRED_COV_TYPE 

433 ) 

434 else: 

435 self.covariance: torch.Tensor = self.compute_covariance( 

436 cov_type=self.cov_type 

437 ) 

438 

439 @icontract.require(lambda self: len(self.support_embeddings) > 0) 

440 def compute_global_moments(self) -> None: 

441 """Method to calculate the global moments of the support embeddings for use in OOD score generation""" 

442 embeddings = torch.cat(list(self.support_embeddings.values())) 

443 self.global_covariance = torch.unsqueeze( 

444 self.compute_covariance_by_type(OOD_COV_TYPE, embeddings), dim=0 

445 ) 

446 global_reg_input = OrderedDict().fromkeys([0], torch.Tensor()) 

447 global_reg_input[0] = self.global_covariance 

448 self.global_covariance: torch.Tensor = self.regularize_covariance( 

449 global_reg_input, OOD_COV_TYPE, "epsilon" 

450 )[0] 

451 self.global_mean: torch.Tensor = torch.mean(embeddings, dim=0) 

452 

453 

454############################################### 

455@beartype 

456class EquineProtonet(Equine): 

457 """ 

458 A class representing an EQUINE model that utilizes protonets and (optionally) relative Mahalanobis distances 

459 to generate OOD and model confidence scores. This wraps any pytorch embedding neural network 

460 and provides the `forward`, `predict`, `save`, and `load` methods required by Equine. 

461 """ 

462 

463 def __init__( 

464 self, 

465 embedding_model: torch.nn.Module, 

466 emb_out_dim: int, 

467 cov_type: CovType = CovType.UNIT, 

468 relative_mahal: bool = True, 

469 use_temperature: bool = False, 

470 init_temperature: float = 1.0, 

471 device: str = "cpu", 

472 feature_names: Optional[list[str]] = None, 

473 label_names: Optional[list[str]] = None, 

474 ) -> None: 

475 """ 

476 EquineProtonet class constructor 

477 

478 Parameters 

479 ---------- 

480 embedding_model : torch.nn.Module 

481 Neural Network feature embedding model. 

482 emb_out_dim : int 

483 The number of output features from the embedding model. 

484 cov_type : CovType, optional 

485 The type of covariance to use when training the protonet [UNIT, DIAG, FULL], by default CovType.UNIT. 

486 relative_mahal : bool, optional 

487 Use relative mahalanobis distance for OOD calculations. If false, uses standard mahalanobis distance instead, by default True. 

488 use_temperature : bool, optional 

489 Whether to use temperature scaling after training, by default False. 

490 init_temperature : float, optional 

491 What to use as the initial temperature (1.0 has no effect), by default 1.0. 

492 device : str, optional 

493 The device to train the equine model on (defaults to cpu). 

494 feature_names : list[str], optional 

495 List of strings of the names of the tabular features (ex ["duration", "fiat_mean", ...]) 

496 label_names : list[str], optional 

497 List of strings of the names of the labels (ex ["streaming", "voip", ...]) 

498 """ 

499 super().__init__( 

500 embedding_model, 

501 device=device, 

502 feature_names=feature_names, 

503 label_names=label_names, 

504 ) 

505 self.cov_type = cov_type 

506 self.cov_reg_type = COV_REG_TYPE 

507 self.relative_mahal = relative_mahal 

508 self.emb_out_dim = emb_out_dim 

509 self.epsilon = DEFAULT_EPSILON 

510 self.outlier_score_kde: OrderedDict[int, gaussian_kde] = OrderedDict() 

511 self.model_summary: dict[str, Any] = dict() 

512 self.use_temperature = use_temperature 

513 self.init_temperature = init_temperature 

514 self.register_buffer( 

515 "temperature", torch.Tensor(self.init_temperature * torch.ones(1)) 

516 ) 

517 

518 self.model: torch.nn.Module = Protonet( 

519 embedding_model, 

520 self.emb_out_dim, 

521 self.cov_type, 

522 self.cov_reg_type, 

523 self.epsilon, 

524 device=device, 

525 ) 

526 

527 def forward(self, X: torch.Tensor) -> torch.Tensor: 

528 """ 

529 Generates logits for classification based on the input tensor. 

530 

531 Parameters 

532 ---------- 

533 X : torch.Tensor 

534 The input tensor for generating predictions. 

535 

536 Returns 

537 ------- 

538 torch.Tensor 

539 The output class predictions. 

540 """ 

541 preds, _ = self.model(X) 

542 return preds 

543 

544 @icontract.require(lambda calib_frac: calib_frac > 0 and calib_frac < 1) 

545 def train_model( 

546 self, 

547 dataset: TensorDataset, 

548 num_episodes: int, 

549 calib_frac: float = 0.2, 

550 support_size: int = 25, 

551 way: int = 3, 

552 episode_size: int = 100, 

553 loss_fn: Callable = torch.nn.functional.cross_entropy, 

554 opt_class: Callable = torch.optim.Adam, 

555 num_calibration_epochs: int = 2, 

556 calibration_lr: float = 0.01, 

557 ) -> dict[str, Any]: 

558 """ 

559 Train or fine-tune an EquineProtonet model. 

560 

561 Parameters 

562 ---------- 

563 dataset : TensorDataset 

564 Input pytorch TensorDataset of training data for model. 

565 num_episodes : int 

566 The desired number of episodes to use for training. 

567 calib_frac : float, optional 

568 Fraction of given training data to reserve for model calibration, by default 0.2. 

569 support_size : int, optional 

570 Number of support examples to generate for each class, by default 25. 

571 way : int, optional 

572 Number of classes to train on per episode, by default 3. 

573 episode_size : int, optional 

574 Number of examples to use per episode, by default 100. 

575 loss_fn : Callable, optional 

576 A pytorch loss function, eg., torch.nn.CrossEntropyLoss(), by default torch.nn.functional.cross_entropy. 

577 opt_class : Callable, optional 

578 A pytorch optimizer, e.g., torch.optim.Adam, by default torch.optim.Adam. 

579 num_calibration_epochs : int, optional 

580 The desired number of epochs to use for temperature scaling, by default 2. 

581 calibration_lr : float, optional 

582 Learning rate for temperature scaling, by default 0.01. 

583 

584 Returns 

585 ------- 

586 tuple[dict[str, Any], torch.Tensor, torch.Tensor] 

587 A tuple containing the model summary, the held out calibration data, and the calibration labels. 

588 """ 

589 self.train() 

590 

591 if self.use_temperature: 

592 self.temperature: torch.Tensor = torch.Tensor( 

593 self.init_temperature * torch.ones(1) 

594 ).type_as(self.temperature) 

595 

596 X, Y = dataset[:] 

597 

598 self.validate_feature_label_names(X.shape[-1], torch.unique(Y).shape[0]) 

599 

600 train_x, calib_x, train_y, calib_y = stratified_train_test_split( 

601 X, Y, test_size=calib_frac 

602 ) 

603 optimizer = opt_class(self.parameters()) 

604 

605 train_x.to(self.device) 

606 train_y.to(self.device) 

607 calib_x.to(self.device) 

608 calib_y.to(self.device) 

609 

610 for i in tqdm(range(num_episodes)): 

611 optimizer.zero_grad() 

612 

613 support, episode_x, episode_y = generate_episode( 

614 train_x, train_y, support_size, way, episode_size 

615 ) 

616 self.model.update_support(support) 

617 

618 _, dists = self.model(episode_x) 

619 loss_value = loss_fn( 

620 torch.neg(dists).to(self.device), episode_y.to(self.device) 

621 ) 

622 loss_value.backward() 

623 optimizer.step() 

624 

625 self.eval() 

626 full_support = generate_support( 

627 train_x, 

628 train_y, 

629 support_size, 

630 selected_labels=torch.unique(train_y).tolist(), 

631 ) 

632 

633 self.model.update_support( 

634 full_support 

635 ) # update support with final selected examples 

636 

637 X_embed = self.model.compute_embeddings(calib_x) 

638 pred_probs, dists = self.model(calib_x) 

639 ood_dists = self._compute_ood_dist(X_embed, pred_probs, dists) 

640 self._fit_outlier_scores(ood_dists, calib_y) 

641 

642 if self.use_temperature: 

643 self.calibrate_temperature( 

644 calib_x, calib_y, num_calibration_epochs, calibration_lr 

645 ) 

646 

647 date_trained = datetime.now().strftime("%m/%d/%Y, %H:%M:%S") 

648 self.train_summary: dict[str, Any] = generate_train_summary( 

649 self, train_y, date_trained 

650 ) 

651 return_dict: dict[str, Any] = dict() 

652 return_dict["train_summary"] = self.train_summary 

653 return_dict["calib_x"] = calib_x 

654 return_dict["calib_y"] = calib_y 

655 return return_dict 

656 

657 def calibrate_temperature( 

658 self, 

659 calib_x: torch.Tensor, 

660 calib_y: torch.Tensor, 

661 num_calibration_epochs: int = 1, 

662 calibration_lr: float = 0.01, 

663 ) -> None: 

664 """ 

665 Fine-tune the temperature after training. Note that this function is also run at the conclusion of train_model. 

666 

667 Parameters 

668 ---------- 

669 calib_x : torch.Tensor 

670 Training data to be used for temperature calibration. 

671 calib_y : torch.Tensor 

672 Labels corresponding to `calib_x`. 

673 num_calibration_epochs : int, optional 

674 Number of epochs to tune temperature, by default 1. 

675 calibration_lr : float, optional 

676 Learning rate for temperature optimization, by default 0.01. 

677 

678 Returns 

679 ------- 

680 None 

681 """ 

682 self.temperature.requires_grad = True 

683 optimizer = torch.optim.Adam([self.temperature], lr=calibration_lr) 

684 for t in range(num_calibration_epochs): 

685 optimizer.zero_grad() 

686 with torch.no_grad(): 

687 pred_probs, dists = self.model(calib_x) 

688 dists = dists.to(self.device) / self.temperature.to(self.device) 

689 loss = torch.nn.functional.cross_entropy( 

690 torch.neg(dists).to(self.device), calib_y.to(torch.long).to(self.device) 

691 ) 

692 loss.backward() 

693 optimizer.step() 

694 self.temperature.requires_grad = False 

695 

696 @icontract.ensure(lambda self: len(self.model.support_embeddings) > 0) 

697 def _fit_outlier_scores( 

698 self, ood_dists: torch.Tensor, calib_y: torch.Tensor 

699 ) -> None: 

700 """ 

701 Private function to fit outlier scores with a kernel density estimate (KDE). 

702 

703 Parameters 

704 ---------- 

705 ood_dists : torch.Tensor 

706 Tensor of computed OOD distances. 

707 calib_y : torch.Tensor 

708 Tensor of class labels for `ood_dists` examples. 

709 

710 Returns 

711 ------- 

712 None 

713 """ 

714 for label in self.model.support_embeddings.keys(): 

715 class_ood_dists = ood_dists[calib_y == int(label)].cpu().detach().numpy() 

716 class_kde = gaussian_kde(class_ood_dists) # TODO convert to torch func 

717 self.outlier_score_kde[label] = class_kde 

718 

719 def _compute_outlier_scores(self, ood_dists, predictions) -> torch.Tensor: 

720 """ 

721 Private function to compute OOD scores using the calculated kernel density estimate (KDE). 

722 

723 Parameters 

724 ---------- 

725 ood_dists : torch.Tensor 

726 Tensor of computed OOD distances. 

727 predictions : torch.Tensor 

728 Tensor of model protonet predictions. 

729 

730 Returns 

731 ------- 

732 torch.Tensor 

733 Tensor of OOD scores for the given examples. 

734 """ 

735 ood_scores = torch.zeros_like(ood_dists) 

736 for i in range(len(predictions)): 

737 # Use KDE and RMD corresponding to the predicted class 

738 predicted_class = int(torch.argmax(predictions[i, :])) 

739 p_value = self.outlier_score_kde[int(predicted_class)].integrate_box_1d( 

740 ood_dists[i].detach().cpu().numpy(), np.inf 

741 ) 

742 ood_scores[i] = 1.0 - np.clip(p_value, 0.0, 1.0) 

743 

744 return ood_scores 

745 

746 @icontract.ensure(lambda result: len(result) > 0) 

747 def _compute_ood_dist( 

748 self, 

749 X_embeddings: torch.Tensor, 

750 predictions: torch.Tensor, 

751 distances: torch.Tensor, 

752 ) -> torch.Tensor: 

753 """ 

754 Private function to compute OOD distances using a distance function. 

755 

756 Parameters 

757 ---------- 

758 X_embeddings : torch.Tensor 

759 Tensor of example embeddings. 

760 predictions : torch.Tensor 

761 Tensor of model protonet predictions for the given embeddings. 

762 distances : torch.Tensor 

763 Tensor of calculated protonet distances for the given embeddings. 

764 

765 Returns 

766 ------- 

767 torch.Tensor 

768 Tensor of OOD distances for the given embeddings. 

769 """ 

770 preds = torch.argmax(predictions, dim=1) 

771 preds = preds.unsqueeze(dim=-1) 

772 # Calculate (Relative) Mahalanobis Distance: 

773 if self.relative_mahal: 

774 null_distance = self.model.compute_distance( 

775 X_embeddings, self.model.global_mean, self.model.global_covariance 

776 ) 

777 null_distance = null_distance.unsqueeze(dim=-1) 

778 ood_dist = distances.gather(1, preds) - null_distance 

779 else: 

780 ood_dist = distances.gather(1, preds) 

781 

782 ood_dist = torch.reshape(ood_dist, (-1,)) 

783 return ood_dist 

784 

785 def predict(self, X: torch.Tensor) -> EquineOutput: 

786 """Predict function for EquineProtonet, inherited and implemented from Equine. 

787 

788 Parameters 

789 ---------- 

790 X : torch.Tensor 

791 Input tensor. 

792 

793 Returns 

794 ------- 

795 EquineOutput 

796 Output object containing prediction probabilities and OOD scores. 

797 """ 

798 X_embed = self.model.compute_embeddings(X) 

799 if X_embed.shape == torch.Size([self.model.emb_out_dim]): 

800 X_embed = X_embed.unsqueeze(dim=0) # Handle single examples 

801 preds, dists = self.model(X) 

802 if self.use_temperature: 

803 dists = dists / self.temperature 

804 preds = torch.softmax(torch.negative(dists), dim=1) 

805 ood_dist = self._compute_ood_dist(X_embed, preds, dists) 

806 ood_scores = self._compute_outlier_scores(ood_dist, preds) 

807 

808 self.validate_feature_label_names(X.shape[-1], preds.shape[-1]) 

809 

810 return EquineOutput(classes=preds, ood_scores=ood_scores, embeddings=X_embed) 

811 

812 @icontract.require(lambda calib_frac: (calib_frac > 0.0) and (calib_frac < 1.0)) 

813 def update_support( 

814 self, 

815 support_x: torch.Tensor, 

816 support_y: torch.Tensor, 

817 calib_frac: float, 

818 label_names: Optional[list[str]] = None, 

819 ) -> None: 

820 """Function to update protonet support examples with given examples. 

821 

822 Parameters 

823 ---------- 

824 support_x : torch.Tensor 

825 Tensor containing support examples for protonet. 

826 support_y : torch.Tensor 

827 Tensor containing labels for given support examples. 

828 calib_frac : float 

829 Fraction of given support data to use for OOD calibration. 

830 label_names : list[str], optional 

831 List of strings of the names of the labels (ex ["streaming", "voip", ...]) 

832 

833 Returns 

834 ------- 

835 None 

836 """ 

837 

838 support_x, calib_x, support_y, calib_y = stratified_train_test_split( 

839 support_x, support_y, test_size=calib_frac 

840 ) 

841 labels, counts = torch.unique(support_y, return_counts=True) 

842 if label_names is not None: 842 ↛ 843line 842 didn't jump to line 843 because the condition on line 842 was never true

843 self.label_names = label_names 

844 self.validate_feature_label_names(support_x.shape[-1], labels.shape[0]) 

845 

846 support = OrderedDict() 

847 for label, count in list(zip(labels.tolist(), counts.tolist())): 

848 class_support = generate_support( 

849 support_x, 

850 support_y, 

851 support_size=count, 

852 selected_labels=[label], 

853 ) 

854 support.update(class_support) 

855 

856 self.model.update_support(support) 

857 

858 X_embed = self.model.compute_embeddings(calib_x) 

859 preds, dists = self.model(calib_x) 

860 ood_dists = self._compute_ood_dist(X_embed, preds, dists) 

861 

862 self._fit_outlier_scores(ood_dists, calib_y) 

863 

864 @icontract.require(lambda self: len(self.model.support) > 0) 

865 def get_support(self) -> OrderedDict[int, torch.Tensor]: 

866 """ 

867 Get the support examples for the model. 

868 

869 Returns 

870 ------- 

871 OrderedDict[int, torch.Tensor] 

872 The support examples for the model. 

873 """ 

874 return self.model.support 

875 

876 @icontract.require(lambda self: len(self.model.prototypes) > 0)