How to return prediction tensors?

Hello.

I have built and trained a model, and I want to test it. The testing I want to do is on un-labeled data, so I only want the actual predictions (tensors), no test loss, or any metric.

In the model class, I have the following function for testing:

    def test_step(self, batch, batch_idx):
        input_ids = batch['input_ids']
        attention_mask = batch['attention_mask']

        outputs = self(input_ids, attention_mask)
        print(outputs)
        return outputs

Somewhere else in my code where I want to get the results, I have:

    def evaluate_model(self) -> None:
        trainer = pl.Trainer(gpus=1)
        predictions = trainer.test(model=self.model,
                                               datamodule=self.data_module)
        print(predictions)

The data module looks like this:


class GenresTestDataModule(pl.LightningDataModule):
    def __init__(self,
                 x_test,
                 tokenizer,
                 max_len: int,
                 batch_size: int
                 ):
        super(GenresTestDataModule, self).__init__()
        self.x_test = x_test
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.batch_size = batch_size
        self.test_dataset = None

    def setup(self, stage: Optional[str] = None) -> None:
        self.test_dataset = GenresDataset(text=self.x_test,
                                          labels=None,
                                          tokenizer=self.tokenizer,
                                          max_length=self.max_len)

    def test_dataloader(self) -> EVAL_DATALOADERS:
        return DataLoader(self.test_dataset,
                          batch_size=self.batch_size,
                          shuffle=True)

The print(outputs) in the test step prints the correct tensors (i.e.
[ 7.2523e-01, -2.4189e-01, -1.7315e-01, -5.9857e-01, -8.6459e-02,
-7.6301e-01, 4.5665e-01, 1.5330e-01, -1.9930e-01, 1.2902e-01,
-9.5539e-02, -1.3968e-01, 1.9320e-01, -4.3536e-02, 2.8427e-01,
1.9324e-01, -6.5592e-01, -5.7927e-03, 1.1902e-01]).

My main issue is that I do not know how to return them from test_step to the evaluate_model funciton.

Thanks!

You can just call your model on the inputs, like you do to get the outputs in test_step.

Thanks, I got it solved in the meanwhile.

I was supposed to use predict_step, in which I would just call the model, and then in the service or wherever a Trainer is instantiated, call trainer.predict(module, datamodule).