Confused logit callback (#118)

* r

* r

* r

* r

* r

* r

* r

* r

* r

* r

* r

* r

* r

* r

* r

* r

* r

* r

* r

* r

* flake8

* plotting

* added more dm tests

* added more dm tests

Co-authored-by: Jirka <jirka@pytorchlightning.ai>

docs/source/index.rst

@@ -21,6 +21,7 @@ PyTorch-Lightning-Bolts documentation
  callbacks
  info_callbacks
  variational_callbacks
+ vision_callbacks
 
 .. toctree::
  :maxdepth: 2

docs/source/variational_callbacks.rst

@@ -11,5 +11,11 @@ Latent Dim Interpolator
 -----------------------
 Interpolates latent dims.
 
+Example output:
+
+ .. image:: _images/gans/basic_gan_interpolate.jpg
+ :width: 400
+ :alt: Example latent space interpolation
+
 .. autoclass:: pl_bolts.callbacks.variational.LatentDimInterpolator
  :noindex:

docs/source/vision_callbacks.rst

@@ -0,0 +1,22 @@
+.. role:: hidden
+ :class: hidden-section
+
+Vision Callbacks
+================
+Useful callbacks for vision models
+
+---------------
+
+Confused Logit
+--------------
+Shows how the input would have to change to move the prediction from one logit to the other
+
+
+Example outputs:
+
+ .. image:: _images/vision/confused_logit.png
+ :width: 400
+ :alt: Example of prediction confused between 5 and 8
+
+.. autoclass:: pl_bolts.callbacks.vision.confused_logit.ConfusedLogitCallback
+ :noindex:

pl_bolts/callbacks/vision/__init__.py

@@ -0,0 +1 @@
+from pl_bolts.callbacks.vision.confused_logit import ConfusedLogitCallback

pl_bolts/callbacks/vision/confused_logit.py

@@ -0,0 +1,127 @@
+import torch
+from torch import nn
+from pytorch_lightning import Callback
+
+
+class ConfusedLogitCallback(Callback): # pragma: no-cover
+
+ def __init__(
+ self,
+ top_k,
+ projection_factor=3,
+ min_logit_value=5.0,
+ logging_batch_interval=20,
+ max_logit_difference=0.1
+ ):
+ """
+ Takes the logit predictions of a model and when the probabilities of two classes are very close, the model
+ doesn't have high certainty that it should pick one vs the other class.
+
+ This callback shows how the input would have to change to swing the model from one label prediction
+ to the other.
+
+ In this case, the network predicts a 5... but gives almost equal probability to an 8.
+ The images show what about the original 5 would have to change to make it more like a 5 or more like an 8.
+
+ For each confused logit the confused images are generated by taking the gradient from a logit wrt an input
+ for the top two closest logits.
+
+ Example::
+
+ from pl_bolts.callbacks.vision import ConfusedLogitCallback
+ trainer = Trainer(callbacks=[ConfusedLogitCallback()])
+
+
+ .. note:: whenever called, this model will look for self.last_batch and self.last_logits in the LightningModule
+
+ .. note:: this callback supports tensorboard only right now
+
+ Args:
+ top_k: How many "offending" images we should plot
+ projection_factor: How much to multiply the input image to make it look more like this logit label
+ min_logit_value: Only consider logit values above this threshold
+ logging_batch_interval: how frequently to inspect/potentially plot something
+ max_logit_difference: when the top 2 logits are within this threshold we consider them confused
+
+ Authored by:
+
+ - Alfredo Canziani
+
+ """
+ super().__init__()
+ self.top_k = top_k
+ self.projection_factor = projection_factor
+ self.max_logit_difference = max_logit_difference
+ self.logging_batch_interval = logging_batch_interval
+ self.min_logit_value = min_logit_value
+
+ def on_batch_end(self, trainer, pl_module):
+
+ # show images only every 20 batches
+ if (trainer.batch_idx + 1) % self.logging_batch_interval != 0:
+ return
+
+ # pick the last batch and logits
+ x, y = pl_module.last_batch
+ logits = pl_module.last_logits
+
+ # only check when it has opinions (ie: the logit > 5)
+ if logits.max() > self.min_logit_value:
+ # pick the top two confused probs
+ (values, idxs) = torch.topk(logits, k=2, dim=1)
+
+ # care about only the ones that are at most eps close to each other
+ eps = self.max_logit_difference
+ mask = (values[:, 0] - values[:, 1]).abs() < eps
+
+ if mask.sum() > 0:
+ # pull out the ones we care about
+ confusing_x = x[mask, ...]
+ confusing_y = y[mask]
+
+ mask_idxs = idxs[mask]
+
+ self._plot(confusing_x, confusing_y, trainer, pl_module, mask_idxs)
+
+ def _plot(self, confusing_x, confusing_y, trainer, model, mask_idxs):
+ from matplotlib import pyplot as plt
+
+ batch_size, c, w, h = confusing_x.size()
+
+ confusing_x = confusing_x[:self.top_k]
+ confusing_y = confusing_y[:self.top_k]
+
+ model.eval()
+ x_param_a = nn.Parameter(confusing_x)
+ x_param_b = nn.Parameter(confusing_x)
+
+ for logit_i, x_param in enumerate((x_param_a, x_param_b)):
+ logits = model(x_param.view(batch_size, -1))
+ logits[:, mask_idxs[:, logit_i]].sum().backward()
+
+ # reshape grads
+ grad_a = x_param_a.grad.view(batch_size, w, h)
+ grad_b = x_param_b.grad.view(batch_size, w, h)
+
+ for img_i in range(len(confusing_x)):
+ x = confusing_x[img_i].squeeze(0)
+ y = confusing_y[img_i]
+ ga = grad_a[img_i]
+ gb = grad_b[img_i]
+
+ mask_idx = mask_idxs[img_i]
+
+ fig, axarr = plt.subplots(nrows=2, ncols=3, figsize=(15, 10))
+ self.__draw_sample(fig, axarr, 0, 0, x, f'True: {y}')
+ self.__draw_sample(fig, axarr, 0, 1, ga, f'd{mask_idx[0]}-logit/dx')
+ self.__draw_sample(fig, axarr, 0, 2, gb, f'd{mask_idx[1]}-logit/dx')
+ self.__draw_sample(fig, axarr, 1, 1, ga * 2 + x, f'd{mask_idx[0]}-logit/dx')
+ self.__draw_sample(fig, axarr, 1, 2, gb * 2 + x, f'd{mask_idx[1]}-logit/dx')
+
+ trainer.logger.experiment.add_figure('confusing_imgs', fig, global_step=trainer.global_step)
+
+ @staticmethod
+ def __draw_sample(fig, axarr, row_idx, col_idx, img, title):
+ im = axarr[row_idx, col_idx].imshow(img)
+ fig.colorbar(im, ax=axarr[row_idx, col_idx])
+ axarr[row_idx, col_idx].set_title(title, fontsize=20)

setup.cfg

@@ -16,6 +16,7 @@ exclude_lines =
  add_model_specific_args
 
 [coverage:run]
+# TODO, remove this ignores in future
 omit =
  pl_bolts/datamodules/stl10_datamodule.py
  pl_bolts/datamodules/ssl_imagenet_datamodule.py
@@ -36,11 +37,13 @@ doctests = True
 verbose = 2
 # https://pep8.readthedocs.io/en/latest/intro.html#error-codes
 format = pylint
+# see: https://www.flake8rules.com/
 ignore =
- E731
- W504
- F401
- F841
+ E731 # Do not assign a lambda expression, use a def
+ W504 # Line break occurred after a binary operator
+ F401 # Module imported but unused
+ F841 # Local variable name is assigned to but never used
+ W605 # Invalid escape sequence 'x'
 
 # setup.cfg or tox.ini
 [check-manifest]