How to get vectors after running model with test data?

[sudarshan85]

Hello,
I'm using the Huggingface library for a text classification project. During testing, I would like access to the [CLS] vectors of the input. Currently, this is the code for my test steps inside my Lightning module:

  @torch.no_grad()
  def test_epoch_end(self, outputs):
    loss = torch.stack(list(zip(*outputs))[0])    
    logits = torch.cat(list(zip(*outputs))[1])    
    preds = logits.argmax(axis=1).cpu()
    labels = torch.stack(list(zip(*outputs))[2]).view(logits.shape[0]).to(torch.int).cpu()   
    cls_vectors = torch.stack(list(zip(*outputs))[3]).view(logits.shape[0], -1).cpu()
    self.log('test_loss', loss)
    self.log('accuracy', accuracy_score(labels, preds))
    self.log('precision', precision_score(labels, preds))
    self.log('recall', recall_score(labels, preds))
    self.log('f1', f1_score(labels, preds))
    self.log('cls', cls_vectors)

  @torch.no_grad()
  def test_step(self, batch, batch_idx):
    outputs = self(**batch, output_hidden_states=True)    
    labels = batch['labels']
    loss = outputs[0]
    logits = outputs[1]    
    cls_vectors = outputs[2][-1][:,0,:]
    return loss, logits, labels, cls_vectors

This is how I run the test:

test_trainer = pl.Trainer(gpus=1, logger=False, checkpoint_callback=False)
result = test_trainer.test(clf_model, dataloaders=test_dl)

Now, the [CLS] vectors are correctly getting accumulated in test_epoch_end method. However, I don't know how to return them from there. I used self.log to log the vectors, however, this seems to log a single float value. Specfically result[0][cls] is a single value instead of the (# samples, 768) dimension tensor. How would I return the [CLS] vectors after running through the test set?

Thanks.

[ShaneTian]

Same issue as you.
I want to return some outputs which are not numerical such as text. But I cannot return them correctly.

---

I noticed that the return value of Trainer.test in the doc is only the recorded metric.
https://pytorch-lightning.readthedocs.io/en/1.6.0/common/trainer.html#test

But it seems that the log can contain a Tensor. Does it not work for you?
https://pytorch-lightning.readthedocs.io/en/1.6.0/common/lightning_module.html#log

[sudarshan85]

But it seems that the log can contain a Tensor. Does it not work for you?

I used log to log the vectors in this line self.log('cls', cls_vectors):

• Nothing got logged to disk. In fact nothing gets logged during testing at all
• It was added to the result dictionary
• Only a single float value was recorded instead of a tensor

Please check my original post for details.

Thanks.

[sudarshan85]

I figured out the problem with logging not working with testing. I didn't enable it! Once enabled logging of metrics works, however, logging of the vectors still only saves a single float. I have a workaround for now. I just save the vectors into the log directory separately and load them later for analysis. Here is the relevant modification:

  @torch.no_grad()
  def test_epoch_end(self, outputs):
    loss = torch.stack(list(zip(*outputs))[0])    
    logits = torch.cat(list(zip(*outputs))[1])    
    preds = logits.argmax(axis=1).cpu().numpy()
    labels = torch.stack(list(zip(*outputs))[2]).view(logits.shape[0]).to(torch.int).cpu().numpy()
    cls_vectors = torch.stack(list(zip(*outputs))[3]).view(logits.shape[0], -1).cpu().numpy()
    with open(f'{self.logger.log_dir}/cls_vectors.npy', 'wb') as f:
      np.save(f, cls_vectors)
    self.log('test_loss', loss, logger=True)
    self.log('accuracy', accuracy_score(labels, preds), logger=True)
    self.log('precision', precision_score(labels, preds), logger=True)
    self.log('recall', recall_score(labels, preds), logger=True)
    self.log('f1', f1_score(labels, preds), logger=True)

Another issue with logging is that both during training and testing the logs are written to metrics.csv inside the same directory if I pass in the same directory. And its not appended, its re-written. This means that when testing after training I lose all of the training metrics. In the current project I don't really use the metrics so its OK, but what I would like is to pass in the name of the metrics (instead of the default metrics.csv) so that I can save the training and testing metrics in separate files (but in the same directory).

[rohitgr7]

to gather and return the predictions/outputs from LightningModule, consider using trainer.predict API.

[sudarshan85]

to gather and return the predictions/outputs from LightningModule, consider using trainer.predict API.

My question was regarding accessing the vectors of the model not the predictions.

[rohitgr7]

My question was regarding accessing the vectors of the model not the predictions.

what do you mean by vectors of the model? whatever you return from predict_step will be accumulated and returned when you call trainer.predict

[sudarshan85]

Quoting from my OP

During testing, I would like access to the [CLS] vectors of the input.

I believe predict only gives me the predictions of the task that I'm working on.

I didn't use the predict_step in my model just the test_epoch_end and I'm able to log my metrics. As I mentioned, I need the 768-dim [CLS] vector (part of the BERT model) for each input when the input runs through the model.

Quoting from one of my replies:

I have a workaround for now. I just save the vectors into the log directory separately and load them later for analysis.

I don't know if the [CLS] vectors can be returned in the predict_step. I'll have to check that out.

[sudarshan85]

My workaround (saving the CLS vectors during test_epoch_end) is working but its not an ideal solution. I want to poke and see if there is a better and easier way to return/log not just metrics (which tend to be scalars) but also vectors/tensors that we are interested in while running the model through the test set.