MisconfigurationException: `training_epoch_end` expects a return of None. HINT: remove the return statement in `training_epoch_end`.

[nkay28]

Bug description

I don;t have a return statement in my training method but i still get the misconfiguration error for not returning none from training_epoch_end. the code breaks aftet he epoch zero. I'm trying to use lightening to train a ResNet (pretrained) model on new storm images data.

How to reproduce the bug

No response

Error messages and logs

---------------------------------------------------------------------------
MisconfigurationException                 Traceback (most recent call last)
Input In [81], in <cell line: 2>()
      1 storm_model = PretrainedWindModel(hparams=hparams)
----> 2 storm_model.fit()

Input In [78], in PretrainedWindModel.fit(self)
    107 def fit(self):
    108     self.trainer = pl.Trainer(
    109         max_epochs=self.max_epochs,
    110         default_root_dir=self.output_path,
   (...)
    123         num_sanity_val_steps=self.hparams.get("val_sanity_checks", 0),
    124     )
--> 125     self.trainer.fit(self)

File C:\Users\Nitin_Kumar\anaconda3\lib\site-packages\pytorch_lightning\trainer\trainer.py:582, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    580     raise TypeError(f"`Trainer.fit()` requires a `LightningModule`, got: {model.__class__.__qualname__}")
    581 self.strategy._lightning_module = model
--> 582 call._call_and_handle_interrupt(
    583     self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    584 )

File C:\Users\Nitin_Kumar\anaconda3\lib\site-packages\pytorch_lightning\trainer\call.py:38, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
     36         return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
     37     else:
---> 38         return trainer_fn(*args, **kwargs)
     40 except _TunerExitException:
     41     trainer._call_teardown_hook()

File C:\Users\Nitin_Kumar\anaconda3\lib\site-packages\pytorch_lightning\trainer\trainer.py:624, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    617 ckpt_path = ckpt_path or self.resume_from_checkpoint
    618 self._ckpt_path = self._checkpoint_connector._set_ckpt_path(
    619     self.state.fn,
    620     ckpt_path,  # type: ignore[arg-type]
    621     model_provided=True,
    622     model_connected=self.lightning_module is not None,
    623 )
--> 624 self._run(model, ckpt_path=self.ckpt_path)
    626 assert self.state.stopped
    627 self.training = False

File C:\Users\Nitin_Kumar\anaconda3\lib\site-packages\pytorch_lightning\trainer\trainer.py:1061, in Trainer._run(self, model, ckpt_path)
   1057 self._checkpoint_connector.restore_training_state()
   1059 self._checkpoint_connector.resume_end()
-> 1061 results = self._run_stage()
   1063 log.detail(f"{self.__class__.__name__}: trainer tearing down")
   1064 self._teardown()

File C:\Users\Nitin_Kumar\anaconda3\lib\site-packages\pytorch_lightning\trainer\trainer.py:1140, in Trainer._run_stage(self)
   1138 if self.predicting:
   1139     return self._run_predict()
-> 1140 self._run_train()

File C:\Users\Nitin_Kumar\anaconda3\lib\site-packages\pytorch_lightning\trainer\trainer.py:1163, in Trainer._run_train(self)
   1160 self.fit_loop.trainer = self
   1162 with torch.autograd.set_detect_anomaly(self._detect_anomaly):
-> 1163     self.fit_loop.run()

File C:\Users\Nitin_Kumar\anaconda3\lib\site-packages\pytorch_lightning\loops\loop.py:200, in Loop.run(self, *args, **kwargs)
    198     self.on_advance_start(*args, **kwargs)
    199     self.advance(*args, **kwargs)
--> 200     self.on_advance_end()
    201     self._restarting = False
    202 except StopIteration:

File C:\Users\Nitin_Kumar\anaconda3\lib\site-packages\pytorch_lightning\loops\fit_loop.py:285, in FitLoop.on_advance_end(self)
    283     epoch_end_outputs = self.trainer._call_lightning_module_hook("training_epoch_end", epoch_end_outputs)
    284     if epoch_end_outputs is not None:
--> 285         raise MisconfigurationException(
    286             "`training_epoch_end` expects a return of None. "
    287             "HINT: remove the return statement in `training_epoch_end`."
    288         )
    289 # free memory
    290 self._outputs = []

MisconfigurationException: `training_epoch_end` expects a return of None. HINT: remove the return statement in `training_epoch_end`.

Environment

Current environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 1.10):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

[carmocca]

Hi! Can you share your LightningModule definition?

[nkay28]

class PretrainedWindModel(pl.LightningModule):
    def __init__(self, hparams):
        super(PretrainedWindModel, self).__init__()

        # self.hparams = hparams
        self.hparams.update(dict(hparams))
        self.learning_rate = self.hparams.get("lr", 2e-4)
        self.hidden_size = self.hparams.get("embedding_dim", 50)
        self.dropout = self.hparams.get("dropout", 0.1)
        self.max_epochs = self.hparams.get("max_epochs", 1)
        self.num_workers = self.hparams.get("num_workers", 0)
        self.batch_size = self.hparams.get("batch_size", 10)
        self.x_train = self.hparams.get("x_train")
        self.y_train = self.hparams.get("y_train")
        self.x_val = self.hparams.get("x_val")
        self.y_val = self.hparams.get("y_val")
        self.num_outputs = 1  # One prediction for regression

        # Where final model will be saved
        self.output_path = Path.cwd() / self.hparams.get("output_path", "model-outputs")
        self.output_path.mkdir(exist_ok=True)

        # Where TensorBoard logs will be saved
        self.log_path = Path.cwd() / self.hparams.get("log_path", "logs")
        self.log_path.mkdir(exist_ok=True)
        # self.logger = pl.loggers.TensorBoardLogger(self.log_path, name="benchmark_model")
        # self.logger = TensorBoardLogger(self.log_path, name="benchmark_model")

        # Instantiate training and validation datasets
        # print(self.x_train.keys())
        self.train_dataset = DatasetWIND(self.x_train, self.y_train)
        # self.train_dataset = WindSpeedImageDataset(self.x_train, self.y_train)
        self.val_dataset = DatasetWIND(self.x_val, self.y_val)
        # self.val_dataset = WindSpeedImageDataset(self.x_val, self.y_val)
        self.model = self.prepare_model()

    def prepare_model(self):
        res_model = models.resnet152(pretrained=True)
        # Input size of 2048 for resnet152
        # https://pytorch.org/hub/pytorch_vision_resnet/
        res_model.fc = nn.Sequential(
            nn.Linear(2048, self.hidden_size),
            nn.ReLU(inplace=True),
            nn.Dropout(self.dropout),
            nn.Linear(self.hidden_size, self.num_outputs),
        )
        return res_model

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

    def training_step(self, batch, batch_idx):
        x = batch["image"]
        y = batch["label"]
        criterion = RMSELoss()
        # Switch to training mode
        loss = criterion(
            self.model.train().forward(x).squeeze(), y.type(torch.FloatTensor)
        )
        tensorboard_logs = {"train_loss": loss}
        return {"loss": loss, "log": tensorboard_logs}

    def validation_step(self, batch, batch_idx):
        x = batch["image"]
        y = batch["label"]
        criterion = RMSELoss()
        # Switch to evaluation mode
        loss = criterion(
            self.model.eval().forward(x).squeeze(), y.type(torch.FloatTensor)
        )
        tensorboard_logs = {"val_loss": loss}
        return {"batch_val_loss": loss, "log": tensorboard_logs}

    def train_dataloader(self):
        return DataLoader(
            self.train_dataset, num_workers=self.num_workers, batch_size=self.batch_size
        )

    def val_dataloader(self):
        return DataLoader(
            self.val_dataset, num_workers=self.num_workers, batch_size=self.batch_size
        )

    def configure_optimizers(self):
        return torch.optim.AdamW(self.model.parameters(), lr=self.learning_rate)

    def training_epoch_end(self, outputs):
        avg_train_loss = torch.stack(tuple(output["loss"] for output in outputs)).mean()
        return {
            "avg_epoch_train_loss": avg_train_loss,
            "progress_bar": {"avg_epoch_train_loss": avg_train_loss},
            "log": {"avg_epoch_train_loss": avg_train_loss},
        }

    def validation_epoch_end(self, outputs):
        avg_val_loss = torch.stack(
            tuple(output["batch_val_loss"] for output in outputs)
        ).mean()
        return {
            "avg_epoch_val_loss": avg_val_loss,
            "progress_bar": {"avg_epoch_val_loss": avg_val_loss},
            "log": {"avg_epoch_val_loss": avg_val_loss},
        }

    ## Convenience Methods ##

    def fit(self):
        self.trainer = pl.Trainer(
            max_epochs=self.max_epochs,
            default_root_dir=self.output_path,
            # logger=self.logger,
            # checkpoint_callback=
            callbacks=[
                pl.callbacks.ModelCheckpoint(
                    # filepath=self.output_path,
                    dirpath=self.output_path,
                    monitor="avg_epoch_val_loss",
                    mode="min",
                    verbose=True,
                )
            ],
            gradient_clip_val=self.hparams.get("gradient_clip_val", 1),
            num_sanity_val_steps=self.hparams.get("val_sanity_checks", 0),
        )
        self.trainer.fit(self)

    @torch.no_grad()
    def make_submission_frame(self, x_test):
        # test_dataset = DatasetWIND(x_test)
        test_dataset = WindSpeedImageDataset(x_test)
        test_dataloader = DataLoader(
            test_dataset, num_workers=self.num_workers, batch_size=self.batch_size
        )
        submission_frame = pd.DataFrame(index=x_test.image_id, columns=["wind_speed"])
        for batch in tqdm(test_dataloader, total=len(test_dataloader)):
            x = batch["image"]
            preds = self.eval().forward(x)
            submission_frame.loc[batch["image_id"], "wind_speed"] = (
                preds.detach().numpy().squeeze()
            )
        submission_frame.wind_speed = submission_frame.wind_speed.astype(float)
        return submission_frame

[nkay28]

If i remove the return statements in def training_epoch_end(self, outputs): and def validation_epoch_end(self, outputs):, I get this error:

MisconfigurationException: `ModelCheckpoint(monitor='avg_epoch_val_loss')` could not find the monitored key in the returned metrics: ['epoch', 'step']. HINT: Did you call `log('avg_epoch_val_loss', value)` in the `LightningModule`?

[carmocca]

@nkay28 The design you are following is from legacy versions of PyTorch Lightning (less than 1.0). You are using multiple patterns that we don't recommend. I strongly suggest you take a dive around our docs to find out how to properly use each feature: https://pytorch-lightning.readthedocs.io/en/stable/

For the error in your last message, you want to do

    def training_epoch_end(self, outputs):
        avg_train_loss = torch.stack(tuple(output["loss"] for output in outputs)).mean()
        self.log("avg_epoch_train_loss": avg_train_loss, logger=True, prog_bar=True)

Without any return

[nkay28]

I see. That's an unfortunate miss on my part. Could you please point to to an example use of models like ResNet, if possible? It would be a great life saver and help me cover for lost time.
I will definitely dig into the docs link you mentioned above. Does this one have a good example to build of off? - https://pytorch-lightning.readthedocs.io/en/stable/notebooks/course_UvA-DL/04-inception-resnet-densenet.html
Sorry, I'm new to using CNNs, or even PyTorch.
Thanks a lot for your feedback!

[carmocca]

We have one in https://github.com/Lightning-AI/lightning/blob/master/examples/pl_domain_templates/imagenet.py

Closing as the original issue is resolved