Save model into separate files
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
I was able to solve my problem by overwriting _del_model and _save_model to perform the desired actions.