TypeError: 'collections.OrderedDict' object is not callable

[neel04]

Hey, training an encoder-decoder model with Lightning ⚡ The task is multi-class image segmentation with image size 256x256

At the end of an epoch, it seems Lightning attempts to save the model for some reason, which yields an error even though its not re-loading that saved checkpoint in any way 🤔

There's no bug report ever filed about this error on S.O, nor Github so I decided to open this discussion 🤗

Code

This is for init and executing trainer,

I know the code snippet is pretty long - but its boilerplate mostly, to allow for some basic mistake I might've made 😃

class Predator(pl.LightningModule):

    def __init__(self, learning_rate, encoder_name, encoder_depth, decoder_channels, out_classes, **kwargs):
        super().__init__()
        
        self.model = Unet(
            encoder_depth=encoder_depth, encoder_name=encoder_name, decoder_channels=decoder_channels, classes=out_classes, 
            encoder_weights='c200k',**kwargs
        ).to(self.device)

        self.model.segmentation_head[1] = torch.nn.ConvTranspose2d(256, 6, kernel_size=(4, 4), stride=(4, 4)).to(self.device)

        self.model.to(self.device)

        self.learning_rate = learning_rate

        # preprocessing parameteres for image
        params = encoders[encoder_name]['pretrained_settings']['Comma200k']
        self.register_buffer("std", torch.tensor(params["std"]).view(1, 3, 1, 1))
        self.register_buffer("mean", torch.tensor(params["mean"]).view(1, 3, 1, 1))

        # for image segmentation dice loss could be the best first choice
        self.loss_fn = FocalLoss(mode='multiclass') #torch.nn.CrossEntropyLoss()

    def forward(self, image):
        # normalize image here
        image = image #(image - self.mean) / self.std
        mask = self.model(image)
        return mask

    def shared_step(self, batch, stage):
        image = batch[0].float()

        h, w = image.shape[2:]
        assert h % 32 == 0 and w % 32 == 0

        tgt_mask = batch[1].squeeze(1).long() #squeezing

        # Shape of the mask should be [batch_size, num_classes, height, width]
        # for binary segmentation num_classes = 1        
        logits_mask = self.forward(image)

        # Predicted mask contains logits, and loss_fn param `from_logits` is set to True

        loss = self.loss_fn(logits_mask, tgt_mask)

        return {
            "loss": loss
            }

    def shared_epoch_end(self, outputs, stage):
      print('\n------------Epoch End------------\n',stage)

    def training_step(self, batch, batch_idx):
        features = batch[0].float()
        masks = batch[1].squeeze(1).long()

        logits = self.forward(features)

        logs = {}

        loss = self.loss_fn(logits, masks)
        print(loss)
        logs["train_loss"] = loss

        return {"loss": loss, "log": logs}          

    def training_epoch_end(self, outputs):
        return self.shared_epoch_end(outputs, "train")

    def validation_step(self, batch, batch_idx):
        return self.shared_step(batch, "valid")

    def validation_epoch_end(self, outputs):
        return self.shared_epoch_end(outputs, "valid")

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

Predator_model = Predator(encoder_name="myencoder", encoder_depth=3, decoder_channels=[32, 32, 32], 
                          out_classes=256, learning_rate=1e-5)

trainer = pl.Trainer(
    accelerator='gpu',
    devices=1, 
    max_epochs=5,
    auto_lr_find=True,
    precision=16,
    auto_scale_batch_size=True
)

trainer.fit(
    Predator_model, 
    train_dataloaders=training_loader, 
    val_dataloaders=validation_loader,
)

---

Error

---------------------------------------------------------------------------

TypeError                                 Traceback (most recent call last)

[<ipython-input-39-e640dcbba470>](https://localhost:8080/#) in <module>()
     14     Predator_model,
     15     train_dataloaders=training_loader,
---> 16     val_dataloaders=validation_loader,
     17 )

23 frames

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in fit(self, model, train_dataloaders, val_dataloaders, datamodule, train_dataloader, ckpt_path)
    739             train_dataloaders = train_dataloader
    740         self._call_and_handle_interrupt(
--> 741             self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    742         )
    743 

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in _call_and_handle_interrupt(self, trainer_fn, *args, **kwargs)
    683         """
    684         try:
--> 685             return trainer_fn(*args, **kwargs)
    686         # TODO: treat KeyboardInterrupt as BaseException (delete the code below) in v1.7
    687         except KeyboardInterrupt as exception:

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in _fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    775         # TODO: ckpt_path only in v1.7
    776         ckpt_path = ckpt_path or self.resume_from_checkpoint
--> 777         self._run(model, ckpt_path=ckpt_path)
    778 
    779         assert self.state.stopped

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in _run(self, model, ckpt_path)
   1197 
   1198         # dispatch `start_training` or `start_evaluating` or `start_predicting`
-> 1199         self._dispatch()
   1200 
   1201         # plugin will finalized fitting (e.g. ddp_spawn will load trained model)

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in _dispatch(self)
   1277             self.training_type_plugin.start_predicting(self)
   1278         else:
-> 1279             self.training_type_plugin.start_training(self)
   1280 
   1281     def run_stage(self):

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py](https://localhost:8080/#) in start_training(self, trainer)
    200     def start_training(self, trainer: "pl.Trainer") -> None:
    201         # double dispatch to initiate the training loop
--> 202         self._results = trainer.run_stage()
    203 
    204     def start_evaluating(self, trainer: "pl.Trainer") -> None:

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in run_stage(self)
   1287         if self.predicting:
   1288             return self._run_predict()
-> 1289         return self._run_train()
   1290 
   1291     def _pre_training_routine(self):

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in _run_train(self)
   1317         self.fit_loop.trainer = self
   1318         with torch.autograd.set_detect_anomaly(self._detect_anomaly):
-> 1319             self.fit_loop.run()
   1320 
   1321     def _run_evaluate(self) -> _EVALUATE_OUTPUT:

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/base.py](https://localhost:8080/#) in run(self, *args, **kwargs)
    143             try:
    144                 self.on_advance_start(*args, **kwargs)
--> 145                 self.advance(*args, **kwargs)
    146                 self.on_advance_end()
    147                 self.restarting = False

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/fit_loop.py](https://localhost:8080/#) in advance(self)
    232 
    233         with self.trainer.profiler.profile("run_training_epoch"):
--> 234             self.epoch_loop.run(data_fetcher)
    235 
    236             # the global step is manually decreased here due to backwards compatibility with existing loggers

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/base.py](https://localhost:8080/#) in run(self, *args, **kwargs)
    149                 break
    150 
--> 151         output = self.on_run_end()
    152         return output
    153 

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py](https://localhost:8080/#) in on_run_end(self)
    296 
    297         # call train epoch end hooks
--> 298         self.trainer.call_hook("on_train_epoch_end")
    299         self.trainer.call_hook("on_epoch_end")
    300         self.trainer.logger_connector.on_epoch_end()

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in call_hook(self, hook_name, pl_module, *args, **kwargs)
   1493             callback_fx = getattr(self, hook_name, None)
   1494             if callable(callback_fx):
-> 1495                 callback_fx(*args, **kwargs)
   1496 
   1497             # next call hook in lightningModule

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/callback_hook.py](https://localhost:8080/#) in on_train_epoch_end(self)
     91         """Called when the epoch ends."""
     92         for callback in self.callbacks:
---> 93             callback.on_train_epoch_end(self, self.lightning_module)
     94 
     95     def on_validation_epoch_start(self):

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py](https://localhost:8080/#) in on_train_epoch_end(self, trainer, pl_module)
    319             and (trainer.current_epoch + 1) % self._every_n_epochs == 0
    320         ):
--> 321             self.save_checkpoint(trainer)
    322         trainer.fit_loop.global_step += 1
    323 

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py](https://localhost:8080/#) in save_checkpoint(self, trainer)
    396         self._save_top_k_checkpoint(trainer, monitor_candidates)
    397         # Mode 2: save monitor=None checkpoints
--> 398         self._save_none_monitor_checkpoint(trainer, monitor_candidates)
    399         # Mode 3: save last checkpoints
    400         self._save_last_checkpoint(trainer, monitor_candidates)

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py](https://localhost:8080/#) in _save_none_monitor_checkpoint(self, trainer, monitor_candidates)
    694         # set the best model path before saving because it will be part of the state.
    695         previous, self.best_model_path = self.best_model_path, filepath
--> 696         trainer.save_checkpoint(filepath, self.save_weights_only)
    697         if self.save_top_k == 1 and previous and previous != filepath:
    698             trainer.training_type_plugin.remove_checkpoint(previous)

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in save_checkpoint(self, filepath, weights_only)
   1911 
   1912     def save_checkpoint(self, filepath: _PATH, weights_only: bool = False) -> None:
-> 1913         self.checkpoint_connector.save_checkpoint(filepath, weights_only)
   1914 
   1915     """

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py](https://localhost:8080/#) in save_checkpoint(self, filepath, weights_only)
    475             weights_only: saving model weights only
    476         """
--> 477         _checkpoint = self.dump_checkpoint(weights_only)
    478         self.trainer.training_type_plugin.save_checkpoint(_checkpoint, filepath)
    479 

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py](https://localhost:8080/#) in dump_checkpoint(self, weights_only)
    377             "global_step": global_step,
    378             "pytorch-lightning_version": pl.__version__,
--> 379             "state_dict": self._get_lightning_module_state_dict(),
    380         }
    381         if _fault_tolerant_training():

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py](https://localhost:8080/#) in _get_lightning_module_state_dict(self)
    489             metric.sync()
    490 
--> 491         state_dict = self.trainer.training_type_plugin.lightning_module_state_dict()
    492 
    493         for metric in metrics:

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py](https://localhost:8080/#) in lightning_module_state_dict(self)
    282         """Returns model state."""
    283         model = self.lightning_module
--> 284         return model.state_dict()
    285 
    286     def save_checkpoint(self, checkpoint: Dict[str, Any], filepath: _PATH) -> None:

[/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in state_dict(self, destination, prefix, keep_vars)
   1257         for name, module in self._modules.items():
   1258             if module is not None:
-> 1259                 module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
   1260         for hook in self._state_dict_hooks.values():
   1261             hook_result = hook(self, destination, prefix, local_metadata)

[/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in state_dict(self, destination, prefix, keep_vars)
   1257         for name, module in self._modules.items():
   1258             if module is not None:
-> 1259                 module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
   1260         for hook in self._state_dict_hooks.values():
   1261             hook_result = hook(self, destination, prefix, local_metadata)

TypeError: 'collections.OrderedDict' object is not callable