Can you print the model summary at an arbitrary time during training?

[mattwarkentin]

Hi,

The model summary (and soon rich model summary) only shows up once at the very beginning of the trainer.fit() procedure when trainer.is_global_zero, but I am wondering if there is a straightforward way for the model summary to be printed at some other point during training?

My use case is that I am doing some transfer learning, where, at the beginning of the fitting procedure, some number of layers are "frozen", and then based on either (1) a certain number of epochs, or (2) some metric has stopped improving, I will unfreeze layers to start updating them.

I have a simple example of the callback shown here:

class UnfreezeLayers(pl.Callback):
    def __init__(self, before_epoch = -1) -> None:
        super().__init__()
        self.before_epoch = before_epoch

    def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        if pl_module.current_epoch == self.before_epoch:
            pl_module.print("Training all parameters now!")
            for param in pl_module.model.parameters():
                param.requires_grad = True

Instead of printing pl_module.print("Training all parameters now!"), I would like to show the model summary again to see how many trainable parameters there are after unfreezing.

[rohitgr7]

does this work?

from pytorch_lightning.utilities.model_summary import summarize

class UnfreezeLayers(pl.Callback):
    def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"):
        print(summarize(pl_module))

[mattwarkentin]

Doesn't look like you can import summarize

>>> from pytorch_lightning.utilities.model_summary import summarize

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'pytorch_lightning.utilities.model_summary'

[rohitgr7]

I tried on master.. https://github.com/PyTorchLightning/pytorch-lightning/blob/8fcdcb598bbf0fbe6cda92415d1d69eb1e2d0aa0/pytorch_lightning/utilities/model_summary.py#L439

try just print(self.summarize())
which version are you using?

[mattwarkentin]

self.summarize() seems to work. Thanks.