Save predictions in `test_step`
Hey,
I set up an AutoEncoder for image reconstruction using the LightningModule (i.e. my_model). At inference, I currently run pl.Trainer.test() to obtain my test metrics, pl.Trainer.predict() to obtain my image predictions, and then in a 3rd loop save each prediction obtained from the pl.Trainer.predict()-step. This is awfully complicated and ideally I would like to directly save predictions in test_step(). I have the following test_step().
def test_step(self, batch, batch_idx):
    batch_hat = self(batch)
    loss_test = nn.functional.mse_loss(batch_hat, batch)
    self.log("loss_test", loss_test, on_step=True, on_epoch=True, sync_dist=True)
    
    # I want to add something like this (NOT FUNCTIONAL)
    filepath = trainer.datamodule.dataset_test.my_file_names[batch_idx]
    my_save_function(batch_hat, filepath)

    return loss_test
and I call test_step() with the conventional
trainer.test(my_model, my_datamodule)
, where my_datamodule is LightningDataModule and trainer is a pl.Trainer().
Question: How can I access the name of my samples in batch within test_step() that are stored in my_datamodule?
Goal: I want to re-use the file names assigned to my input images for the predictions in test_step().
Best,
dsethz
hey @dsethz !
you can return the filenames from your Dataset.__getitem__ which will be available inside the batch and you can access them inside test_step easily.