Loading model for prediction yields RuntimeError: Error(s) in loading state_dict

[stillsen]

Hi,
I want to use a pretrained ResNet or DenseNet with adjusted fc layer for image regression.
After training, however, when I want to load the finetuned DenseNet for prediction I get a RuntimeError: Error(s) in loading state_dict.
This error does not occur for the ResNet.
My best guess is that this is somehow linked to checkpointing and that the value for self.model_arch might not be properly stored in the state_dict?
But I am not even sure if I need to tell the LightningModule what to save in the state_dict.
Any ideas what might go wrong?

Big Thanks already!

LightningModule constructor:

class ImageRegressor(LightningModule):
    def __init__(self, optimizer:str = 'adam', k:int = 0, lr:float = 1e-3, batch_size:int = 16, pretrained:bool = True, tune_fc_only:bool = False, model: str = 'resnet50'):
        super().__init__()
        self.lr = lr
        self.batch_size = batch_size
        self.k = k
        
        optimizers = {'adam': Adam, 'sgd': SGD}
        self.optimizer = optimizers[optimizer]
        self.criterion = nn.MSELoss(reduction='mean')
        self.model_arch = model
        num_target_classes = 1

        if self.model_arch == 'resnet50':
            # init a pretrained resnet
            self.model = models.resnet50(pretrained=pretrained)
            num_filters = self.model.fc.in_features
            self.model.fc = nn.Sequential(
                nn.ReLU(),
                nn.Linear(num_filters, num_target_classes))
        elif self.model_arch == 'densenet':
            print('setting up densenet')
            self.model = models.densenet121(pretrained=pretrained)
            num_filters = self.model.classifier.in_features
            self.model.classifier = nn.Sequential(
                nn.ReLU(),
                nn.Linear(num_filters, num_target_classes))
        if pretrained:
            if tune_fc_only: # option to only tune the fully-connected layers
                for child in list(self.model.children())[:-1]:
                    for param in child.parameters():
                        param.requires_grad = False

Saving is done automatically using this checkpoint:

ModelCheckpoint(dirpath=this_output_dir,
                                     filename='model_'+str(k)+'_{epoch}.pt',
                                     monitor='val_loss')

loading checkpoint:

model = ImageRegressor(pretrained=True, tune_fc_only=True, model='densenet')
type(model).load_from_checkpoint(path)

[rohitgr7]

should be:

model = ImageRegressor.load_from_checkpoint(path, pretrained=True, tune_fc_only=True, model='densenet')

check out the examples and description here: https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#load-from-checkpoint

slightly similar discussion: #12399

[stillsen]

That does the trick. Thank you so much!