Save model into separate files

[lvoegtlin]

Hei everyone

I am working on a small framework based on PyTorch lightning to perform some experiments.
Our models always are a composition out of two networks (backbone and some header) as we want to be able to replace them independently.

class EncoderHeaderModel(pl.LightningModule):

        def __init__(self, backbone: Union[pl.LightningModule, torch.nn.Module],
                 header: Union[pl.LightningModule, torch.nn.Module]):
        super().__init__()

        # sanity check if the last layer of the backbone is compatible with the first layer of the header

        self.backbone = backbone
        self.header = header

    def forward(self, x):
        x = self.backbone(x)
        x = self.header(x)
        return x

Now we have the problem with saving the state_dict of these two models separately. To store the whole model we are using the model_checkpoint callback which works fine.

Is there an easy way to save the models each time model_checkpoint would save the whole model (I am already experimenting with a subclass of model_checkpoint)? Or should we after the training just extract the state_dicts out of the checkpoint? Any idea how to make this in a "clean" way?

Thanks

[lvoegtlin]

I was able to solve my problem by overwriting _del_model and _save_model to perform the desired actions.

[tchaton]

That's great.
By the way, do you know about Lightning Flash: https://github.com/PyTorchLightning/lightning-flash. It is PyTorch Lightning Team high level framework :)

Best,
T.C