In Masterarbeit:"Anomalie-Detektion in Zellbildern zur Anwendung der Leukämieerkennung" verwendete CSI Methode.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

temperature_scaling.py 4.6KB

2 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  1. import torch
  2. from torch import nn, optim
  3. from torch.nn import functional as F
  4. class ModelWithTemperature(nn.Module):
  5. """
  6. A thin decorator, which wraps a model with temperature scaling
  7. model (nn.Module):
  8. A classification neural network
  9. NB: Output of the neural network should be the classification logits,
  10. NOT the softmax (or log softmax)!
  11. """
  12. def __init__(self, model):
  13. super(ModelWithTemperature, self).__init__()
  14. self.model = model
  15. self.temperature = nn.Parameter(torch.ones(1) * 0.5)
  16. def forward(self, input):
  17. logits = self.model(input)
  18. return self.temperature_scale(logits)
  19. def temperature_scale(self, logits):
  20. """
  21. Perform temperature scaling on logits
  22. """
  23. # Expand temperature to match the size of logits
  24. temperature = self.temperature.unsqueeze(1).expand(logits.size(0), logits.size(1))
  25. return logits / temperature
  26. # This function probably should live outside of this class, but whatever
  27. def set_temperature(self, valid_loader):
  28. """
  29. Tune the tempearature of the model (using the validation set).
  30. We're going to set it to optimize NLL.
  31. valid_loader (DataLoader): validation set loader
  32. """
  33. self.cuda()
  34. nll_criterion = nn.CrossEntropyLoss().cuda()
  35. ece_criterion = _ECELoss().cuda()
  36. # First: collect all the logits and labels for the validation set
  37. logits_list = []
  38. labels_list = []
  39. with torch.no_grad():
  40. for input, label in valid_loader:
  41. input = input.cuda()
  42. logits = self.model(input)
  43. logits_list.append(logits)
  44. labels_list.append(label)
  45. logits = torch.cat(logits_list).cuda()
  46. labels = torch.cat(labels_list).cuda()
  47. # Calculate NLL and ECE before temperature scaling
  48. before_temperature_nll = nll_criterion(logits, labels).item()
  49. before_temperature_ece = ece_criterion(logits, labels).item()
  50. print('Before temperature - NLL: %.3f, ECE: %.3f' % (before_temperature_nll, before_temperature_ece))
  51. # Next: optimize the temperature w.r.t. NLL
  52. optimizer = optim.LBFGS([self.temperature], lr=0.0001, max_iter=50000)
  53. def eval():
  54. loss = nll_criterion(self.temperature_scale(logits), labels)
  55. loss.backward()
  56. return loss
  57. optimizer.step(eval)
  58. # Calculate NLL and ECE after temperature scaling
  59. after_temperature_nll = nll_criterion(self.temperature_scale(logits), labels).item()
  60. after_temperature_ece = ece_criterion(self.temperature_scale(logits), labels).item()
  61. print('Optimal temperature: %.3f' % self.temperature.item())
  62. print('After temperature - NLL: %.3f, ECE: %.3f' % (after_temperature_nll, after_temperature_ece))
  63. return self
  64. class _ECELoss(nn.Module):
  65. """
  66. Calculates the Expected Calibration Error of a model.
  67. (This isn't necessary for temperature scaling, just a cool metric).
  68. The input to this loss is the logits of a model, NOT the softmax scores.
  69. This divides the confidence outputs into equally-sized interval bins.
  70. In each bin, we compute the confidence gap:
  71. bin_gap = | avg_confidence_in_bin - accuracy_in_bin |
  72. We then return a weighted average of the gaps, based on the number
  73. of samples in each bin
  74. See: Naeini, Mahdi Pakdaman, Gregory F. Cooper, and Milos Hauskrecht.
  75. "Obtaining Well Calibrated Probabilities Using Bayesian Binning." AAAI.
  76. 2015.
  77. """
  78. def __init__(self, n_bins=15):
  79. """
  80. n_bins (int): number of confidence interval bins
  81. """
  82. super(_ECELoss, self).__init__()
  83. bin_boundaries = torch.linspace(0, 1, n_bins + 1)
  84. self.bin_lowers = bin_boundaries[:-1]
  85. self.bin_uppers = bin_boundaries[1:]
  86. def forward(self, logits, labels):
  87. softmaxes = F.softmax(logits, dim=1)
  88. confidences, predictions = torch.max(softmaxes, 1)
  89. accuracies = predictions.eq(labels)
  90. ece = torch.zeros(1, device=logits.device)
  91. for bin_lower, bin_upper in zip(self.bin_lowers, self.bin_uppers):
  92. # Calculated |confidence - accuracy| in each bin
  93. in_bin = confidences.gt(bin_lower.item()) * confidences.le(bin_upper.item())
  94. prop_in_bin = in_bin.float().mean()
  95. if prop_in_bin.item() > 0:
  96. accuracy_in_bin = accuracies[in_bin].float().mean()
  97. avg_confidence_in_bin = confidences[in_bin].mean()
  98. ece += torch.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin
  99. return ece