test_dataloader must be implemented to be used with the Lightning Trainer

[Antoine101]

Hi everyone,
I am getting the following error:

Traceback (most recent call last):
  File "d:\Dev\ml_tools\ml_tools\audio\deep_learning\main_trainvaltest.py", line 198, in <module>
    trainer.test(ckpt_path="best")
  File "C:\Users\APU\Anaconda3\lib\site-packages\lightning\pytorch\trainer\trainer.py", line 735, in test
    return call._call_and_handle_interrupt(
  File "C:\Users\APU\Anaconda3\lib\site-packages\lightning\pytorch\trainer\call.py", line 42, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "C:\Users\APU\Anaconda3\lib\site-packages\lightning\pytorch\trainer\trainer.py", line 778, in _test_impl
    results = self._run(model, ckpt_path=ckpt_path)
  File "C:\Users\APU\Anaconda3\lib\site-packages\lightning\pytorch\trainer\trainer.py", line 973, in _run
    results = self._run_stage()
  File "C:\Users\APU\Anaconda3\lib\site-packages\lightning\pytorch\trainer\trainer.py", line 1009, in _run_stage
    return self._evaluation_loop.run()
  File "C:\Users\APU\Anaconda3\lib\site-packages\lightning\pytorch\loops\utilities.py", line 177, in _decorator
    return loop_run(self, *args, **kwargs)
  File "C:\Users\APU\Anaconda3\lib\site-packages\lightning\pytorch\loops\evaluation_loop.py", line 98, in run
    self.setup_data()
  File "C:\Users\APU\Anaconda3\lib\site-packages\lightning\pytorch\loops\evaluation_loop.py", line 150, in setup_data
    dataloaders = _request_dataloader(source)
  File "C:\Users\APU\Anaconda3\lib\site-packages\lightning\pytorch\trainer\connectors\data_connector.py", line 330, in _request_dataloader
    return data_source.dataloader()
  File "C:\Users\APU\Anaconda3\lib\site-packages\lightning\pytorch\trainer\connectors\data_connector.py", line 297, in dataloader
    return call._call_lightning_module_hook(self.instance.trainer, self.name, pl_module=self.instance)
  File "C:\Users\APU\Anaconda3\lib\site-packages\lightning\pytorch\trainer\call.py", line 144, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "C:\Users\APU\Anaconda3\lib\site-packages\lightning\pytorch\core\hooks.py", line 445, in test_dataloader
    raise MisconfigurationException("`test_dataloader` must be implemented to be used with the Lightning Trainer")
lightning.fabric.utilities.exceptions.MisconfigurationException: `test_dataloader` must be implemented to be used with the Lightning Trainer

Although as shown below my test_dataloader is correctly implemented:

class TrainValTestDataModule(LightningDataModule):

    def __init__(self, metadata, dataset_path, batch_size, num_workers, feature_name, feature_processing_parameters, signal_augmentation, feature_augmentation, augmentation_parameters, to_gpus):
        super().__init__()
        self.save_hyperparameters(ignore=["metadata", "dataset_path", "to_gpus"])
        self.prepare_data_per_node = True
        self.metadata = metadata
        self.dataset_path = dataset_path
        self.to_gpus = to_gpus


    def prepare_data(self) -> None:
        return super().prepare_data()
    

    def setup(self, stage=None):
        if stage=="fit":
            self.train_ds = TrainValTestDataset(
                                            metadata=self.metadata, 
                                            dataset_path=self.dataset_path,
                                            stage="train", 
                                            feature_name=self.hparams.feature_name, 
                                            feature_processing_parameters=self.hparams.feature_processing_parameters,
                                            signal_augmentation=self.hparams.signal_augmentation,
                                            feature_augmentation=self.hparams.feature_augmentation,
                                            augmentation_parameters=self.hparams.augmentation_parameters,
                                            to_gpus=self.to_gpus,
                                            devices=self.trainer.strategy.root_device
                                            )
                                            
            self.validation_ds = TrainValTestDataset(
                                            metadata=self.metadata,
                                            dataset_path=self.dataset_path,
                                            stage="validation", 
                                            feature_name=self.hparams.feature_name, 
                                            feature_processing_parameters=self.hparams.feature_processing_parameters, 
                                            signal_augmentation=False,
                                            feature_augmentation=False,
                                            augmentation_parameters=self.hparams.augmentation_parameters,
                                            to_gpus=self.to_gpus,
                                            devices=self.trainer.strategy.root_device
                                            )
        
        elif stage=="test":
            self.test_ds = TrainValTestDataset(
                                            metadata=self.metadata,
                                            dataset_path=self.dataset_path,
                                            stage=stage, 
                                            feature_name=self.hparams.feature_name, 
                                            feature_processing_parameters=self.hparams.feature_processing_parameters, 
                                            signal_augmentation=False,
                                            feature_augmentation=False,
                                            augmentation_parameters=self.hparams.augmentation_parameters,
                                            to_gpus=self.to_gpus,
                                            devices=self.trainer.strategy.root_device
                                            )            
        

    def train_dataloader(self):
        train_dataloader = DataLoader(
                                    dataset=self.train_ds, 
                                    batch_size=self.hparams.batch_size, 
                                    shuffle=True,
                                    num_workers=self.hparams.num_workers, 
                                    pin_memory=True
                                    )
        return train_dataloader


    def val_dataloader(self):
        validation_dataloader = DataLoader(
                                    dataset=self.validation_ds, 
                                    batch_size=self.hparams.batch_size,
                                    shuffle=False, 
                                    num_workers=self.hparams.num_workers, 
                                    pin_memory=True
                                    )
        return validation_dataloader
    
    
    def test_dataloader(self):
        test_dataloader = DataLoader(
                                    dataset=self.test_ds, 
                                    batch_size=self.hparams.batch_size,
                                    shuffle=False, 
                                    num_workers=self.hparams.num_workers, 
                                    pin_memory=True
                                    )
        return test_dataloader

I switched from pytorch-lightning to lightning 2.0.6 as i appeared as a solution in a couple of online forum threads, but it didn't work for me. Downgrading to lightning 2.0.0 didn't fix anything neither.

Here is my LightningModule:

class Module(LightningModule):
    def __init__(self, n_classes, classes_map, optimizer, optimizer_parameters, lr_scheduler, lr_scheduler_parameters, learning_rate, batch_size, model):
        super().__init__()
        # Save hyperparameters to the checkpoint
        self.save_hyperparameters(ignore=["model"])   
        self.model = model     
        #self.loss = nn.BCEWithLogitsLoss()
        self.loss = nn.CrossEntropyLoss()
        # Instantiation of the metrics
        self.train_accuracy = Accuracy(task="multiclass", num_classes=self.hparams.n_classes, average="weighted")
        self.val_accuracy = Accuracy(task="multiclass", num_classes=self.hparams.n_classes, average="weighted")
        self.val_precision = Precision(task="multiclass", num_classes=self.hparams.n_classes, average="weighted")
        self.val_recall = Recall(task="multiclass", num_classes=self.hparams.n_classes, average="weighted")
        self.val_f1_score = F1Score(task="multiclass", num_classes=self.hparams.n_classes, average="weighted")          
        self.val_confmat = ConfusionMatrix(task="multiclass", num_classes=self.hparams.n_classes, normalize="true")
        self.test_accuracy = Accuracy(task="multiclass", num_classes=self.hparams.n_classes, average="weighted")
        self.test_precision = Precision(task="multiclass", num_classes=self.hparams.n_classes, average="weighted")
        self.test_recall = Recall(task="multiclass", num_classes=self.hparams.n_classes, average="weighted")
        self.test_f1_score = F1Score(task="multiclass", num_classes=self.hparams.n_classes, average="weighted")          
        self.test_confmat = ConfusionMatrix(task="multiclass", num_classes=self.hparams.n_classes, normalize="true")


    def configure_optimizers(self): 
        if self.hparams.optimizer == "Adam": 
            optimizer = Adam(self.parameters(), lr=self.hparams.learning_rate)

        if self.hparams.lr_scheduler == "ReduceLROnPlateau":
            scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, patience=self.hparams.lr_scheduler_parameters["patience"], verbose=True)  

        return {
                "optimizer": optimizer,
                "lr_scheduler": {
                        "scheduler": scheduler,
                        "monitor": "validation_loss",
                        "frequency": 1
                        }
                }
    
        
    def training_step(self, batch, batch_idx): 
        index, audio_name, targets, inputs = batch
        logits = self.model(inputs) 
        loss = self.loss(logits, targets)
        predictions = torch.argmax(logits, dim=1)
        self.train_accuracy(logits, targets)
        self.log("training_loss", loss, on_step=False, on_epoch=True, batch_size=self.hparams.batch_size, prog_bar=True)
        self.log("training_accuracy", self.train_accuracy, on_step=False, on_epoch=True, batch_size=self.hparams.batch_size)        
        return {"inputs":inputs, "targets":targets, "predictions":predictions, "loss":loss}
    
    
    def train_epoch_end(self, outputs):
        # Log weights and biases for all layers of the model
        for name, params in self.named_parameters():
            self.logger.experiment.add_histogram(name, params,self.current_epoch)
        # Only after the first training epoch, log one of the training inputs as a figure and log the model graph
        if self.current_epoch == 0:
            input_sample = outputs[0]["inputs"][0]
            input_sample_target = outputs[0]["targets"][0].item()
            input_sample_class = self.hparams.classes_map[input_sample_target]
            fig = plt.figure(figsize=(20,20))
            ax = fig.add_subplot(111)
            ax.imshow(torch.squeeze(input_sample).cpu(), cmap="viridis", origin="lower", aspect="auto")
            ax.set_title(f"Class: {input_sample_class}")
            ax.set_xlabel("Time Frames")
            self.logger.experiment.add_figure("Training sample input", fig)
            input_sample = torch.unsqueeze(input_sample, 3)
            input_sample = torch.permute(input_sample, (3,0,1,2))
            self.logger.experiment.add_graph(self.model, input_sample)

            
    def validation_step(self, batch, batch_idx):
        index, audio_name, targets, inputs = batch
        logits = self.model(inputs)
        loss = self.loss(logits, targets)
        predictions = torch.argmax(logits, dim=1)
        self.val_accuracy(predictions, targets)
        self.val_precision(predictions, targets)
        self.val_recall(predictions, targets)
        self.val_f1_score(predictions, targets)
        self.val_confmat.update(predictions, targets)
        self.log(f"val_loss", loss, on_step=False, on_epoch=True, batch_size=self.hparams.batch_size, prog_bar=True)
        self.log(f"val_accuracy", self.val_accuracy, on_step=False, on_epoch=True, batch_size=self.hparams.batch_size)
        self.log(f"val_precision", self.val_precision, on_step=False, on_epoch=True, batch_size=self.hparams.batch_size)
        self.log(f"val_recall", self.val_recall, on_step=False, on_epoch=True, batch_size=self.hparams.batch_size)
        self.log(f"val_f1_score", self.val_f1_score, on_step=False, on_epoch=True, batch_size=self.hparams.batch_size)        

    
    def on_validation_epoch_end(self):
        # Compute the confusion matrix, turn it into a DataFrame, generate the plot and log it
        cm = self.val_confmat.compute()
        cm = cm.cpu()
        self.val_confmat.reset()
        df_cm = pd.DataFrame(cm.numpy(), index=range(self.hparams.n_classes), columns=range(self.hparams.n_classes))
        fig, ax = plt.subplots(figsize=(16,16))
        sns.heatmap(data=df_cm, annot=True, cmap="Blues", ax=ax, vmin=0, vmax=1)
        ax.set_xticklabels(labels=list(self.hparams.classes_map.values()))
        ax.set_yticklabels(labels=list(self.hparams.classes_map.values()))
        ax.tick_params(axis="y", labelrotation=0)
        ax.tick_params(axis="x", labelrotation=90)
        ax.set_xlabel("Predicted", fontsize=12)
        ax.set_ylabel("Real", fontsize=12)
        self.logger.experiment.add_figure("Confusion matrix", fig, self.current_epoch)
        
    def on_save_checkpoint(self, checkpoint):
        # Get the state_dict from self.model to get rid of the "model." prefix
        checkpoint["state_dict"] = self.state_dict()

    def test_step(self, batch, batch_idx):
        index, audio_name, targets, inputs = batch
        logits = self.model(inputs)
        loss = self.loss(logits, targets)
        predictions = torch.argmax(logits, dim=1)
        self.test_accuracy(predictions, targets)
        self.test_precision(predictions, targets)
        self.test_recall(predictions, targets)
        self.test_f1_score(predictions, targets)
        self.test_confmat.update(predictions, targets)
        self.log(f"test_loss", loss, on_step=False, on_epoch=True, batch_size=self.hparams.batch_size, prog_bar=True)
        self.log(f"test_accuracy", self.test_accuracy, on_step=False, on_epoch=True, batch_size=self.hparams.batch_size)
        self.log(f"test_precision", self.test_precision, on_step=False, on_epoch=True, batch_size=self.hparams.batch_size)
        self.log(f"test_recall", self.test_recall, on_step=False, on_epoch=True, batch_size=self.hparams.batch_size)
        self.log(f"test_f1_score", self.test_f1_score, on_step=False, on_epoch=True, batch_size=self.hparams.batch_size)

[changspencer]

I'm new to PyTorch Lightning, so I have a simple question from what I've quickly skimmed on their documentation/guide for validating and testing a model. Thanks for the post; I'm learning more ways to modularize my old PyTorch code by reading your post and going to the guide.

Where do you tell the Trainer what LightningDataModule (LDM) to use? As far as I can tell, you don't give the Trainer a reference to a LDM in the provided code. Is there some place you pass it in before the trainer.fit(...) method call?

[jxtngx]

we would declare a datamodule and pass it to the trainer with:

import lightning as L

from your_custom_neural_net import MyCustomModel
from your_custom_datamodule import MyCustomDataModule

trainer = L.Trainer()
model = MyCustomModel()
datamodule = MyCustomDataModule()
trainer.fit(model, datamodule)

and an example custom LightningDataModule may look like:

import lightning.pytorch as pl
from torch.utils.data import random_split, DataLoader

# Note - you must have torchvision installed for this example
from torchvision.datasets import MNIST
from torchvision import transforms


class MNISTDataModule(pl.LightningDataModule):
    def __init__(self, data_dir: str = "./"):
        super().__init__()
        self.data_dir = data_dir
        self.transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

    def prepare_data(self):
        # download
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage: str):
        # Assign train/val datasets for use in dataloaders
        if stage == "fit":
            mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)
            self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])

        # Assign test dataset for use in dataloader(s)
        if stage == "test":
            self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)

        if stage == "predict":
            self.mnist_predict = MNIST(self.data_dir, train=False, transform=self.transform)

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=32)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=32)

    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=32)

    def predict_dataloader(self):
        return DataLoader(self.mnist_predict, batch_size=32)

and here is some more reading on LightningDataModules -> https://lightning.ai/docs/pytorch/stable/data/datamodule.html

[Antoine101]

Hi @changspencer and @JustinGoheen ,

Thanks for your inputs!

Stupid mistake from my end... As you said @changspencer, I don't pass the datamodule to my Trainer when calling the .test() method.

My main code before the fix:

# Instantiation of the trainer
    trainer = Trainer(
        accelerator=config['run']['accelerator'],
        devices=config['run']['devices'],
        max_epochs=int(config['run']['epochs']),
        logger=tensorboard_logger,
        log_every_n_steps=1,
        callbacks=[early_stopping, lr_monitor, checkpoint, progress_bar]
    )

    # Train the model
    trainer.fit(model=lm, datamodule=dm) 

    # Evaluate it on the test set
    trainer.test(ckpt_path="best")

My main code after the fix:

    # Instantiation of the trainer
    trainer = Trainer(
        accelerator=config['run']['accelerator'],
        devices=config['run']['devices'],
        max_epochs=int(config['run']['epochs']),
        logger=tensorboard_logger,
        log_every_n_steps=1,
        callbacks=[early_stopping, lr_monitor, checkpoint, progress_bar]
    )

    # Train the model
    trainer.fit(model=lm, datamodule=dm) 

    # Evaluate it on the test set
    trainer.test(ckpt_path="best", datamodule=dm)

Giving the datamodule to the test() method solved my error.

I got mislead by this part of the lightning documentation:

There is no way to pass a datamodule or model to the Trainer at instantiation. So I don't see how it would work doing it like the doc says.

[jxtngx]

can you please share a link to the docs page that led to the mistake.

[Antoine101]

Here it is:
(https://lightning.ai/docs/pytorch/stable/common/evaluation_intermediate.html)

The example, a snapshot of which I have included in my previous reply, is what mislead me, as it only specifies the checkpoint but no datamodule.

[jxtngx]

thank you for the link. I'll be addressing issues like this in #18327. I've prioritized that docs page.

[subhashnerella]

The documentation is misleading. I ran into the same issue.

[coffeered]

In fact, the documentation is also accurate and well-founded.

Within PyTorch Lightning, there exist multiple avenues for accessing DataLoaders.
One approach involves the periodic passing of a DataModule when initiating various stages, such as fit, evaluate, test, or predict.
Alternatively, another method is equally effective, whereby the DataLoaders are tightly integrated with your model, a step accomplished by passing the model to the Trainer during its initialization.

But yes, it can be confusing at times