"Error: 'MyDataModule' object is not iterable" when calling trainer.predict(model, data)

[dianemarquette]

Hi!

I've trained a model successfully. I now want to have a look at the model predictions. I've overridden the predict_dataloader method in my DataModule:

def predict_dataloader(self):
        pred_loader = torch.utils.data.DataLoader(
            self.val_ds, batch_size=1, num_workers=4)
        return pred_loader

Then, I initialize the model using my checkpoint and call the predict method:

checkpoint_dir = os.path.join(root_dir, "logs2/epoch=199-val_loss=0.26-val_dice=1.67.ckpt")

# initialize the data module 
data = KeriDataModule(data_dir=data_dir, pix_dim=(0.6, 0.6, 0.937))

# initialize the LightningModule (from checkpoint)
net = Net.load_from_checkpoint(checkpoint_path=checkpoint_dir)

# initialize Lightning's trainer
trainer = pl.Trainer(gpus=[0])

results = trainer.predict(net, data)

Unfortunately, I get the following error

Predicting: 0it [43:22, ?it/s]

GPU available: True, used: True
TPU available: False, using: 0 TPU cores
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]


Predicting: 0it [00:00, ?it/s]

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-12-4bd0b31315bd> in <module>
     10 trainer = pl.Trainer(gpus=[0])
     11 
---> 12 results = trainer.predict(net, data)
     13 print(results)
     14 

~/miniconda3/envs/monai-lightning-latest/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in predict(self, model, dataloaders, datamodule, return_predictions)
    629         self.data_connector.attach_data(model, predict_dataloaders=dataloaders, datamodule=datamodule)
    630 
--> 631         results = self._run(model)
    632 
    633         assert self.state.stopped

~/miniconda3/envs/monai-lightning-latest/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in _run(self, model)
    754 
    755         # dispatch `start_training` or `start_evaluating` or `start_predicting`
--> 756         self.dispatch()
    757 
    758         # plugin will finalized fitting (e.g. ddp_spawn will load trained model)

~/miniconda3/envs/monai-lightning-latest/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in dispatch(self)
    793             self.accelerator.start_evaluating(self)
    794         elif self.predicting:
--> 795             self.accelerator.start_predicting(self)
    796         else:
    797             self.accelerator.start_training(self)

~/miniconda3/envs/monai-lightning-latest/lib/python3.7/site-packages/pytorch_lightning/accelerators/accelerator.py in start_predicting(self, trainer)
    100 
    101     def start_predicting(self, trainer: 'pl.Trainer') -> None:
--> 102         self.training_type_plugin.start_predicting(trainer)
    103 
    104     def pre_dispatch(self, trainer: 'pl.Trainer') -> None:

~/miniconda3/envs/monai-lightning-latest/lib/python3.7/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py in start_predicting(self, trainer)
    150     def start_predicting(self, trainer: 'pl.Trainer') -> None:
    151         # double dispatch to initiate the predicting loop
--> 152         self._results = trainer.run_stage()
    153 
    154     def training_step(self, *args, **kwargs):

~/miniconda3/envs/monai-lightning-latest/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in run_stage(self)
    804             return self.run_evaluate()
    805         if self.predicting:
--> 806             return self.run_predict()
    807         return self.run_train()
    808 

~/miniconda3/envs/monai-lightning-latest/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py in run_predict(self)
   1071             dataloader = self.accelerator.process_dataloader(dataloader)
   1072             dl_max_batches = self.predict_loop.max_batches[dataloader_idx]
-> 1073             for batch_idx, batch in enumerate(dataloader):
   1074                 if batch is None:
   1075                     continue

TypeError: 'KeriDataModule' object is not iterable

I don't understand what I messed up ^^'. Any help / tip would be greatly appreciated :)

[carmocca]

Sorry! We did not add support for trainer.predict(model, datamodule). We'll do it asap!

You need to do trainer.predict(model, datamodule=datamodule)

[dianemarquette]

Thanks, trainer.predict(model=model, datamodule=datamodule) solved the issue :)