How to use predict function to return predictions

[talhaanwarch]

How to get predictions

  class OurModel(LightningModule):
      def __init__(self):
        super(OurModel,self).__init__()
        self.layer = MyModelV3()
      def forward(self,x):
        return self.layer(x)

  def train_dataloader(self):
    return DataLoader(DataReader(train_df))

  def training_step(self,batch,batch_idx):
    return loss

  def test_dataloader(self):
    return DataLoader(DataReader(test_df))
    
  def test_step(self,batch,batch_idx):
    image,label=batch
    out=self(image)
    loss=self.criterion(out,label)
    return loss

  def predict(self, batch):
        return self(batch) 

I am not sure, how to use predict function. How to define data loader for predict function. I want to get predictions for test_df. But now idea how to do this.

[ethanwharris]

Hi @talhaanwarch, in order to get predictions from a data loader you need to implement predict_step in your LightningModule (docs here: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#predict-step). You would then be able to call Trainer.predict with the dataloader you want use following the API here: https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.html#pytorch_lightning.trainer.trainer.Trainer.predict

Hope that helps 😃

[talhaanwarch]

@ethanwharris
i have already checked predict_step but i am not sure how data loader pass to it. Becuase i am not using datamodule

in doc it is something link this

def predict_step(self, batch, batch_idx: int , dataloader_idx: int = None):
        return self(batch)
model = OurModel()
trainer.predict(model, datamodule)

But my data loader is in OurModel class. So how can I pass data loader to it

[ethanwharris]

Hi @talhaanwarch you would pass your dataloader to the Trainer.predict method. So pseudocode would look something like this:

MyModule(pl.LightningModule):
    ...
    def predict_step(self, batch):
        inputs, targets = batch
        return my_model(inputs)

predict_dataloader = ...

trainer = Trainer(MyModule(...), ...)
predictions = trainer.predict(predict_dataloader)

Note that depending on your use case you may not need to implement the predict_step as it just calls forward by default.

Hope that helps 😃

[talhaanwarch]

what I got I need to initialize my test_loader out of MyModule class again?

[ethanwharris]

@talhaanwarch You have a few options. You could:

• implement a predict_dataloader in your lightning module (like you do with train / test) to give a special dataloader for predicting
• implement predict_dataloader to just return self.test_dataloader
• just call Trainer.predict(my_model, my_model.test_dataloader())
• create a predict dataloader outside of your model and pass it to Trainer.predict

Any of those options should work as expected.

[tchaton]

Dear @talhaanwarch,

To create a dataloder, just do DataLoader(predict_dataset).

predict_dataloader = DataLoader(predict_dataset)

class Model(LightningModule):

    def __init__(self, ...):
        super().__init__()
        self.save_hyperparameters()
        self.model = ...

    def forward(self, batch):
        return self.model(batch)

    def predict_step(self, batch, batch_idx, dataloder_idx = None):
        return self(batch)

trainer = Trainer(...)
predictions = trainer.predict(model, dataloaders=predict_dataloader)