Save predictions in test_step

[dsethz]

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

[rohitgr7]

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.

[dsethz]

Hey @rohitgr7 ,

Thank you for your feedback.

Is there also a way to access my_datamodule and its attributes directly within test_step() (maybe through trainer)?

[rohitgr7]

yes, using self.trainer.datamodule

[dsethz]

Perfect, just what I was looking for.

Thank you!