Optimizing Layer Conductance

ajhepburn: <p>So I’m currently at odds with how to approach this. At first, I was only really using Layer Conductance on a document-by-document basis. However, I’m now trying to scale this a little bit as I’m using the attributions towards a paper. The set that I’m passing in for analysis of an already trained model, consists of 6658 documents. This is taking me, on average, 1 to 1.5 hours to process on a Titan RTX. So for around 120 models, it’s taking me just shy of about 3 days to process it if I split this across 4-6 Titan RTXs (I have GPUs available to me via my work cluster so it grabs how many are available).</p>
<p>The process I am currently following is the BERT SQUAD tutorial, slightly adapted for classification. I am using HuggingFace’s BERT Model as a base with dropout + linear layer on top for output.</p>
<p>I was wondering what optimization strategies were available to me, given my situation. I’ve checked the input tensors and they’re float32. I’ve noticed executing cpu().detach().numpy() slows things down a fair bit. Should I maybe be running this on a batch of documents at a time to avoid running the whole embedding process over and over?</p>
<p>I’ve attached my eval script for reference:</p>
<pre><code class="lang-auto">def run(docs):
    res = []
    for i, doc in enumerate(docs):
        print(f"Processing Doc #{i+1}/{len(docs)}")
        text_ids = tokenizer.encode(doc, add_special_tokens=False)
        input_ids = [tokenizer.cls_token_id] + text_ids + [tokenizer.sep_token_id]
        ref_input_ids = [tokenizer.cls_token_id] + [tokenizer.pad_token_id] * len(text_ids) + [tokenizer.sep_token_id]
        input_ids = torch.tensor([input_ids], device=DEVICE)
        ref_input_ids = torch.tensor([ref_input_ids], device=DEVICE)

        sep_id = len(text_ids)
        seq_len = input_ids.size(1)
        token_type_ids = torch.tensor(
            [[0 if i &lt;= sep_id else 1 for i in range(seq_len)]],
            device=DEVICE
        )
        ref_token_type_ids = torch.zeros_like(token_type_ids, device=DEVICE)

        seq_length = input_ids.size(1)
        position_ids = torch.arange(seq_length, dtype=torch.long, device=DEVICE)
        position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
        ref_position_ids = torch.zeros(seq_length, dtype=torch.long, device=DEVICE)
        ref_position_ids = ref_position_ids.unsqueeze(0).expand_as(input_ids)
        attention_mask = torch.ones_like(input_ids)

        indices = input_ids[0].detach().tolist()
        all_tokens = tokenizer.convert_ids_to_tokens(indices)
        
        input_embeds = model.base.embeddings(input_ids, token_type_ids=token_type_ids, position_ids=position_ids)
        ref_input_embeds = model.base.embeddings(ref_input_ids, token_type_ids=token_type_ids, position_ids=position_ids)
        
        def forward(input_embeds, attention_mask):
            logits = model(inputs_embeds=input_embeds, attention_mask=attention_mask)
            return torch.softmax(logits, dim=1)
    
        def summarize_attributions(attributions):
            attributions = attributions.sum(dim=-1).squeeze(0)
            attributions = attributions / torch.norm(attributions)
            return attributions
        
        layer_attrs = []
        for i in range(model.base.config.num_hidden_layers):
            lc = LayerConductance(forward, model.base.encoder.layer[i])
            layer_attributions = lc.attribute(
                inputs=input_embeds,
                baselines=ref_input_embeds,
                additional_forward_args=(attention_mask),
                target=1
            )
            layer_attrs.append(summarize_attributions(layer_attributions).numpy())
        
        res.append ({
            "doc_no": i,
            "tokens": all_tokens,
            "conductance": layer_attrs
        })
        torch.cuda.empty_cache()
    return res
</code></pre>

Narine: <p>Hi <a class="mention" href="/u/ajhepburn">@ajhepburn</a>, if you run layer conductance on a batch of documents it can definitely speed up things. Only you want to make sure that you have sufficient memory and perhaps create sub-batches combing multiple documents if you can’t fit everything into a memory. Note that the internal requirements of layer conductance is <code>n_steps</code> * <span class="hashtag">#examples</span>. n_steps is defaulted to 50 and you can change it if necessary.</p>
<p>As I can see, you are treating each document as one example, is that right ? If your documents aren’t too large you can definitely try to combine them.<br>
Another option is to use DataParallels and GPUs.<br>
More info about Dataparallels and Distributed Dataparallels can be found here:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://captum.ai/tutorials/Distributed_Attribution">
  <header class="source">
      <img src="https://captum.ai/img/captum.ico" class="site-icon" width="256" height="256">

      <a href="https://captum.ai/tutorials/Distributed_Attribution" target="_blank" rel="noopener nofollow ugc">captum.ai</a>
  </header>

  <article class="onebox-body">
    <img src="https://captum.ai/img/captum-icon.png" class="thumbnail onebox-avatar" width="500" height="500">

<h3><a href="https://captum.ai/tutorials/Distributed_Attribution" target="_blank" rel="noopener nofollow ugc">Captum · Model Interpretability for PyTorch</a></h3>

  <p>Model Interpretability for PyTorch</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>