[Python] Calculating the partial derivative of PyTorch model output with respect to neurons...

I am working on neuron importance for ANNs (in a classification setting). A simple baseline is the partial derivative of the model output for the correct class with respect to the given neuron pre-activation. My goal is to get important neurons for a cluster of related inputs. Thus, I am currently just averaging the importances for the samples in each cluster to get the per-cluster importances. The problem is that I get very small partial derivatives for each neuron on average and that the resulting importance scores are not good in the sense that if you remove the most important neurons from the model, the performance of the model does not decrease more than when compared to removing neurons at random. Other methods, such as using the pre-activations as importance scores, work better in that regard. However, experiments from other authors suggest that the partial derivative based approach should work reasonably well in many settings. Thus, I suspect there is a bug in my code.

My current setup is as following. I know the code is not optimal from a performance POV - it is written with clarity in mind. My question is, if there is a mistake in the code to calculate the partial derivatives with respect to the neurons' pre-activations. Assume that the provided layer in the attribute function is a PyTorch Linear layer. Thus the last argument in the hook should be the partial derivative with respect to the pre-activation.

class AttributionMethod(ABC):
def __init__(self, model: nn.Module, data_loader: Dict[int, DataLoader]):
self.model = model
self.data_loader = data_loader
if self.model is not None:
self.model.eval()

@abstractmethod
def attribute(self, layer: Module, latent_class: int) -> np.ndarray:
pass


class GradientAttribution(AttributionMethod):
def attribute(self, layer: Module, latent_class: int) -> np.ndarray:
# Calculates the importance scores of the neurons in the provided layer for the provided cluster
device = next(self.model.parameters()).device

inputs, targets = next(iter(self.data_loader[latent_class]))
inputs = inputs.to(device)
targets = targets.to(device)
inputs.requires_grad = True

attributions = []

def hook_fn(module, grad_input, grad_output):
attributions.append(grad_output[0].detach().cpu())

handle = layer.register_full_backward_hook(hook_fn)

# Forward and backward pass
for i in range(len(inputs)):
input_i = inputs.unsqueeze(0) # Add batch dimension
output = self.model(input_i).squeeze()
target_i = targets.item()

output[target_i].backward()
# We do not need to zero the grads since only the model param grads are accumulated

handle.remove()

# Concatenate over batch and average per neuron
grads_tensor = torch.cat(attributions, dim=0)
grads_per_neuron = grads_tensor.mean(dim=0).numpy()
return grads_per_neuron

Continue reading...