load checkpoint model error

[morestart]

this is my model __init__ func

def __init__(self, num_classes: int, image_channels: int = 3, drop_rate: int = 0.5,
                 filter_config: tuple = (64, 128, 256, 512, 512), attention=False):

this is my load code:

m = SegNet(num_classes=1)
model = m.load_from_checkpoint('checkpoints/epoch=99-step=312499.ckpt')

when i try to load a checkpoint model, i got this error:

TypeError: __init__() missing 1 required positional argument: 'num_classes'

who can help me?

[tchaton]

Dear @morestart,

Great question ! You should use save_hyperparameters function, so Lightning can save your init arguments inside the checkpoint for future reload.

Here is the associated doc: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html?highlight=save_hyperparameters#save-hyperparameters

class SegNet(LightningModule)

    def __init__(
        self,
        num_classes: int,
        image_channels: int = 3,
        drop_rate: int = 0.5,
        filter_config: tuple = (64, 128, 256, 512, 512),
        attention=False
    ):
        super().__init__()
        self.save_hyperparameters()

        ....

Best,
T.C

[tchaton]

Small note. Some arguments can't be serialised like DataModule or nn.Module. Therefore, you would need to provide them directly to load_from_checkpoint.

Example:

m.load_from_checkpoint('checkpoints/epoch=99-step=312499.ckpt', model=model)

[morestart]

thanks very much, i will try it~