Loading model for prediction yields RuntimeError: Error(s) in loading state_dict
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)
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