How to implement Trainer.predict to save test images #9278
Unanswered
bibinwils
asked this question in
Lightning Trainer API: Trainer, LightningModule, LightningDataModule
Replies: 1 comment 4 replies
-
You have to implement the predict_step(...) function in your LightningModule. Now inside of this function you can torch.save(...) your outputs (which I assume are going to be segmentations?). In order to have the name of the output file to be the same of the input file, my suggestion would be to include the name of the input file in the batch (if you've defined your own DataLoader), something like this:
|
Beta Was this translation helpful? Give feedback.
4 replies
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Uh oh!
There was an error while loading. Please reload this page.
-
Hi, I trained a UNet with input images and mask. For that, I split the entire dataset to train - valid - test. I have trained the network and got the output metrics for the test dataset. There are around 100 test images. Now I would like to save all the test data outputs to a folder with the same name as the test images. For example, if the test image name is image1, then I need to save the test output image also as image1. I have gone through the PyTorch lightning documentation, I come to know about Trainer.predict, but I don't know how to implement it.
Beta Was this translation helpful? Give feedback.
All reactions