Save Test Predictions

1

Documentation is a bit lacking in this area, atleast what I have tried to look through.
Here is my training_step():

def training_step(self, train_batch, batch_idx):
    X_batch = train_batch[0]
    y_batch = train_batch[1]
    y_hat_batch = self.model(X_batch)
    loss = F.cross_entropy(y_hat_batch, y_batch) # CrossEntropyLoss
    
    #Logs
    self.log('training_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
    self.log('train_acc', self.train_acc(y_hat_batch, y_batch), on_step=True, on_epoch=True, prog_bar=True, logger=True)
    
    return loss

I having a hard time thinking/coding the converstion from pytorch to pytorch lightning for the test loop. Here is the pytorch code that includes what I am wanting to save off.

    # Loop through all test batches
    y_hat = []

    with torch.no_grad():
        model.eval()

        for batch_num, data in enumerate(full_test_dataloader):
            # Get data and pass to cuda if CUDA is available
            X_batch = data[0].cuda() if CUDA else data[0]
            y_batch = data[1].cuda() if CUDA else data[1]

            # Infer from Model
            y_hat_batch = model(X_batch)
            _, y_hat_batch_tags = torch.max(y_hat_batch, dim=-1)
            y_hat.append(y_hat_batch_tags.cpu().numpy())

          
    y_hat = np.array([item for sub_list in y_hat for item in sub_list.squeeze().tolist()])
    print(f"# Actual/Predicted Non-Targets: {len(testing_dataset.Y[testing_dataset.Y == 0])}/{len(y_hat[y_hat == 0])}")
    print(f"# Actual/Predicted Targets: {len(testing_dataset.Y[testing_dataset.Y == 1])}/{len(y_hat[y_hat == 1])}")
          
    # Get Accuracy
    total = len(full_test_pytorch_dataset.Y)
    correct = (full_test_pytorch_dataset.Y == y_hat).sum()
    total_accuracy = correct/total
 
    #Save results
    test_dataset_results = {
        "y": full_test_pytorch_dataset.Y,
        "y_hat": y_hat,
        "total_accuracy": total_accuracy
    }

print("Accuracies:")
for dataset in all_test_datasets_results.keys():
    print(f"\t{dataset}: {all_test_datasets_results[dataset]['total_accuracy']}")

#Plot y (image) and y_hat (image) side by side

I really need to have the same functionality and results as the code above.

What I have so far is:

def test_step(self, test_batch, batch_idx):
    X_batch = train_batch[0]
    y_batch = train_batch[1]
    y_hat_batch = self.model(X_batch)
    _, y_hat_batch_tags = torch.max(y_hat_batch, dim=-1)
    
    #Logs
    self.log('train_acc', self.test_acc(y_hat_batch_tags, y_batch), on_step=False, on_epoch=True, prog_bar=True, logger=True)
   
    return

I really like PyTorch Lightning and really want to use it. Any help would greatly be appreciated.

2

This is my solution so far. I made 2 new classes variable in my LightningModule, self.test_y and self.test_y_hat.

Here is my test_step:

def test_step(self, test_batch, batch_idx):
    X_batch = test_batch[0]
    y_batch = test_batch[1]
    y_hat_batch = self.model(X_batch)
    _, y_hat_batch_tags = torch.max(y_hat_batch, dim=-1)
    
    #Logs
    self.log('test_acc', self.test_acc(y_hat_batch_tags, y_batch), on_step=True, on_epoch=True, prog_bar=True, logger=True)
    
    #Save y and y_hat_batch_tags
    y_batch = y_batch.cpu().numpy()
    self.test_y.extend(y_batch)
    
    y_hat_batch = y_hat_batch_tags.cpu().numpy()
    self.test_y_hat.extend(y_hat_batch)
    
    return

I can get the y and y_hat from the test loop this way.
Seems like there should be a more elegant way to do this…

Thoughts?

3

Yeah, I expect trainer.test() to pass the return value of test_epoch_end(), or outputs if test_epoch_end() is not defined.

Maybe they can add a trainer.predict() to do this. I guess I will create an issue.

1 Like