| 288 | self.fc_layers = fc_layers |
| 289 | |
| 290 | def compute_map(self, x, class_idx=None, layer_idx=-1, **kwargs): # type: ignore[override] |
| 291 | logits, acti, _ = self.nn_module(x, **kwargs) |
| 292 | acti = acti[layer_idx] |
| 293 | if class_idx is None: |
| 294 | class_idx = logits.max(1)[-1] |
| 295 | b, c, *spatial = acti.shape |
| 296 | acti = torch.split(acti.reshape(b, c, -1), 1, dim=2) # make the spatial dims 1D |
| 297 | fc_layers = self.nn_module.get_layer(self.fc_layers) |
| 298 | output = torch.stack([fc_layers(a[..., 0]) for a in acti], dim=2) |
| 299 | output = torch.stack([output[i, b : b + 1] for i, b in enumerate(class_idx)], dim=0) |
| 300 | return output.reshape(b, 1, *spatial) # resume the spatial dims on the selected class |
| 301 | |
| 302 | def __call__(self, x, class_idx=None, layer_idx=-1, **kwargs): # type: ignore[override] |
| 303 | """ |