Is it possible to save a file from test_step() function?

user_15864414 0
Viewed 519

I am trying to implement MNIST digits using PyTorch Lightning.

The train function is like the below one

def train(epochs, train_loader, test_loader, model):
        early_stopping = EarlyStopping('train_loss', mode='min', patience=5)
        model_checkpoint = ModelCheckpoint(dirpath=model_path/'mnist_{epoch}-{train_loss:.2f}',monitor='train_loss', mode='min', save_top_k=3)
        trainer = pl.Trainer(max_epochs=epochs, profiler=False, callbacks = [model_checkpoint],default_root_dir=model_path) 
        trainer.fit(model, train_dataloader=train_loader)
        trainer.test(test_dataloaders=test_loader, ckpt_path=None)

The test_step function is like the below one

def test_step(self, test_batch):
    x, y = test_batch
    logits = self.forward(x)
    loss = self.mean_squared_error_loss(logits.squeeze(-1), y.float())

    # I want to calculate R2, MAPE, etc and want to save in a pandas df and 
    # need to return to the train function

    self.log('test_loss', loss)
    return {'test_loss': loss}

I can do calculate R2, MAPE, etc using TorchMetrics. But, I am not sure how (or is it possible) to save them in a pandas df (or maybe in a list) for the whole test dataset. I have gone through this post but not sure how should I try!

Any suggestions are appreciated.

pytorch-lightning pytorch python-3.x python
1 Answers

You can aggregate test result in test_epoch_end:

def test_step(self, test_batch):
    x, y = test_batch
    logits = self.forward(x)
    loss = self.mean_squared_error_loss(logits.squeeze(-1), y.float())


    self.log('test_loss', loss)
    return {'test_loss': loss, "logits":logits, "labels": y}

def test_epoch_end(self, outputs):
    all_preds, all_labels = [], []
    for output in outputs:
        probs = list(output['logits'].cpu().detach().numpy()) # predicted values
        labels = list(output['labels'].flatten().cpu().detach().numpy())
        all_preds.extend(probs)
        all_labels.extend(labels)

    # you can calculate R2 here or save results as file
    r2 = ...

Note that this only works on a single GPU. If you are using multiple GPUs, you need some function to gather results from different GPUs.

To get model predictions, you need to add a predict_step() in the model class.

def predict_step(self, test_batch):
    x, y = test_batch
    logits = self.forward(x)
    return {'logits': logits, 'labels':y}

And run:

outputs = trainer.predict(model, test_loader, return_predictions=True)
user_109539580
Related