code: def compute_heads_importance(
      args,
      model,
      eval_dataloader,
      compute_entropy=True,
      compute_importance=True,
      head_mask=None,
      actually_pruned=False,
  ):
      """This method shows how to compute:
      - head attention entropy
      - head importance scores according to http://arxiv.org/abs/1905.10650
      """
      # Prepare our tensors
      n_layers, n_heads = (
          model.config.num_hidden_layers,
          model.config.num_attention_heads,
      )
      head_importance = torch.zeros(n_layers, n_heads).to(args.device)
      attn_entropy = torch.zeros(n_layers, n_heads).to(args.device)
  
      if head_mask is None:
          head_mask = torch.ones(n_layers, n_heads).to(args.device)
  
      head_mask.requires_grad_(requires_grad=True)
      # If actually pruned attention multi-head, set head mask to None to avoid shape mismatch
      if actually_pruned:
          head_mask = None
  
      tot_tokens = 0.0
      total_loss = 0.0
      for step, inputs in enumerate(
          tqdm(
              eval_dataloader,
              desc="Iteration",
              disable=args.local_rank not in [-1, 0],
          )
      ):
          inputs = tuple(t.to(args.device) for t in inputs)
          (input_ids,) = inputs
  
          # Do a forward pass (not with torch.no_grad() since we need gradients for importance score - see below)
          outputs = model(
              input_ids, labels=input_ids, head_mask=head_mask
          )
          #  (loss), lm_logits, presents, (all hidden_states), (attentions)
          loss, _, all_attentions = (
              outputs[0],
              outputs[1],
              outputs[-1],
          )  # Loss and logits are the first, attention the last
          loss.backward()  # Backpropagate to populate the gradients in the head mask
          total_loss += loss.detach().cpu().numpy()
          if compute_entropy:
              for layer, attn in enumerate(all_attentions):
                  masked_entropy = entropy(attn.detach(), True)
                  attn_entropy[layer] += (
                      masked_entropy.sum(-1).sum(0).sum(0).detach()
                  )
  
          if compute_importance:
              head_importance += head_mask.grad.abs().detach()
          tot_tokens += (
              torch.ones_like(input_ids).float().detach().sum().data
          )
  
      # Normalize
      attn_entropy /= tot_tokens
      head_importance /= tot_tokens
      # Layerwise importance normalization
      if not args.dont_normalize_importance_by_layer:
          exponent = 2
          norm_by_layer = torch.pow(
              torch.pow(head_importance, exponent).sum(-1), 1 / exponent
          )
          head_importance /= norm_by_layer.unsqueeze(-1) + 1e-20
  
      if not args.dont_normalize_global_importance:
          head_importance = (
              head_importance - head_importance.min()
          ) / (head_importance.max() - head_importance.min())
  
      # Print matrices
      if compute_entropy:
          logger.info("Attention entropies")
          print_2d_tensor(attn_entropy)
      if compute_importance:
          logger.info("Head importance scores")
          print_2d_tensor(head_importance)
      logger.info("Head ranked by importance scores")
      head_ranks = torch.zeros(
          head_importance.numel(), dtype=torch.long, device=args.device
      )
      head_ranks[
          head_importance.view(-1).sort(descending=True)[1]
      ] = torch.arange(head_importance.numel(), device=args.device)
      head_ranks = head_ranks.view_as(head_importance)
      print_2d_tensor(head_ranks)
      return attn_entropy, head_importance, total_loss
  
file_path: transformers/examples/research_projects/bertology/run_prune_gpt.py
imports: import argparse
  import logging
  import os
  from datetime import datetime
  import numpy
  import torch
  from torch import nn
  from torch.utils.data import DataLoader, RandomSampler, TensorDataset
  from tqdm import tqdm
  from transformers import GPT2LMHeadModel
label_desc: Use image node_id 4 for calling a global function with example usage: compute_heads_importance(args, model, eval_dataloader, compute_entropy, compute_importance, head_mask, actually_pruned) and returns: attn_entropy,  head_importance,  total_loss