How to use accuracy with PyTorch Lightning?

Chris Staff asked 6 months ago
1 Answers
Best Answer
Chris Staff answered 6 months ago

It is possible to use accuracy by computing it in the test_step, as follows.

  def test_step(self, batch, batch_idx):
    x, y = batch
    x = x.view(x.size(0), -1)
    y_hat = self.layers(x)
    loss = self.ce(y_hat, y)
    y_hat = torch.argmax(y_hat, dim=1)
    accuracy = torch.sum(y == y_hat).item() / (len(y) * 1.0)
    output = dict({
        'test_loss': loss,
        'test_acc': torch.tensor(accuracy),
    return output

In this step, the following is happening:

  • We decompose the batch into x (inputs) and y (targets).
  • We generate y_hat, or the predictions, by feeding our input through our model (defined in self.layers in __init__, can really be any layer / model).
  • We compute the loss value, then the instances of the y_hats, and finally accuracy.
  • We add everything into a dict and return the output as a whole.

You then get these results when calling model.test(..) in Lightning:

{'test_acc': tensor(.96), 'test_loss': tensor(0.0024, device='cuda:0')}

Your Answer

17 + 8 =