How to properly load checkpoint for testing?

[polars05]

I've trained a system as follows:

model = CoolSystem(hparams)

checkpoint_callback = pl.callbacks.ModelCheckpoint(
          filepath=os.path.join(os.getcwd(), 'checkpoints'),
          verbose=True,
          monitor='val_acc', 
          mode='max', 
          prefix='try',
          save_top_k=-1,
          period=1
      )

trainer = pl.Trainer(
      max_epochs=hparams.epochs,
      checkpoint_callback=checkpoint_callback)  

trainer.fit(model)
trainer.test()

And with the above, the test accuracy is 0.7975

However, when I load the checkpoints separately instead:

model_test = CoolSystem(hyperparams)
model_test.load_from_checkpoint('checkpoints/try_ckpt_epoch_1.ckpt')

trainer = pl.Trainer()  

trainer.test(model_test)

The accuracy returned is 0.5705

Am I loading the checkpoints wrongly?

[awaelchli]

Is this the CoolSystem from the examples or is it something custom? Basically you need to tell us how to reproduce it.
Maybe there is something wrong in how you accumulate the test acc? Do you get the same values if you run
trainer.test(model_test) multiple times (for both cases)?

[polars05]

@awaelchli, here is my CoolSystem for reference:

class CoolSystem(pl.LightningModule):

    def __init__(self, hparams):
        super(CoolSystem, self).__init__()

        self.hparams = hparams
        self.data_dir = self.hparams.data_dir
        self.num_classes = self.hparams.num_classes # len(self.hparams.classes)

        ########## define the model ########## 
        arch = torchvision.models.resnet18(pretrained=True)
        num_ftrs = arch.fc.in_features

        modules = list(arch.children())[:-1] 
        self.backbone = torch.nn.Sequential(*modules) # [bs, 512, 1, 1]
        self.final = torch.nn.Sequential(
               torch.nn.Linear(num_ftrs, 128),
               torch.nn.ReLU(inplace=True),
               torch.nn.Linear(128, self.num_classes),
               torch.nn.Softmax(dim=1))   
        ########## /define the model ########## 

        self.training_acc_across_batches_at_curr_epoch = []

    def forward(self, x):
        x = self.backbone(x)
        x = x.reshape(x.size(0), -1) 
        x = self.final(x)
        
        return x
    
    def configure_optimizers(self):
        # REQUIRED
        # can return multiple optimizers and learning_rate schedulers
        # (LBFGS it is automatically supported, no need for closure function)

        optimizer = torch.optim.SGD(list(self.backbone.parameters()) + list(self.final.parameters()), lr=0.001, momentum=0.9)

        exp_lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
 
        return [optimizer], [exp_lr_scheduler]

    def training_step(self, batch, batch_idx):
        # REQUIRED
        x, y = batch
        y_hat = self.forward(x)
        
        loss = F.cross_entropy(y_hat, y)
        
        _, preds = torch.max(y_hat, dim=1)
        acc = torch.sum(preds == y.data) / (y.shape[0] * 1.0)
        # torch.tensor(acc)
        self.training_acc_across_batches_at_curr_epoch.append(acc.item())

        return {'loss': loss} #, 'train_acc': acc}
    
    def on_epoch_end(self):
        train_acc_mean = np.mean(self.training_acc_across_batches_at_curr_epoch)
        
        self.logger.experiment.add_scalar('train_acc', train_acc_mean, global_step=self.current_epoch)
        self.training_acc_across_batches_per_epoch = []  # reset for next epoch

    def validation_step(self, batch, batch_idx):
        # OPTIONAL
        x, y = batch
        y_hat = self.forward(x)

        _, preds = torch.max(y_hat, 1)
        acc = torch.sum(preds == y.data) / (y.shape[0] * 1.0)

        return {'val_loss': F.cross_entropy(y_hat, y), 'val_acc': acc}

    def validation_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        avg_acc = torch.stack([x['val_acc'].float() for x in outputs]).mean()

        print ("val accuracy at end of epoch: ", avg_acc)

        logs = {'val_loss': avg_loss, 'val_acc': avg_acc}
        return {'avg_val_loss': avg_loss, 'avg_val_acc': avg_acc, 'log': logs, 'progress_bar': logs}

    def test_step(self, batch, batch_idx):
        # OPTIONAL
        x, y = batch
        y_hat = self.forward(x)

        _, preds = torch.max(y_hat, 1)
        acc = torch.sum(preds == y.data) / (y.shape[0] * 1.0)

        return {'test_loss': F.cross_entropy(y_hat, y), 'test_acc': acc}

    def test_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
        avg_acc = torch.stack([x['test_acc'].float() for x in outputs]).mean()

        logs = {'test_loss': avg_loss, 'test_acc': avg_acc}
        
        print ('#### Average accuracy: ', avg_acc)
        
        return {'avg_test_loss': avg_loss, 'avg_test_acc': avg_acc, 'log': logs, 'progress_bar': logs}

    @pl.data_loader
    def train_dataloader(self):
        # REQUIRED

        transform = transforms.Compose([
                                transforms.RandomResizedCrop(224),
                                transforms.RandomHorizontalFlip(),
                                transforms.ToTensor(),
                                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                                ])

        train_set = torchvision.datasets.ImageFolder(os.path.join(self.data_dir, 'train'), transform)
        train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True, num_workers=4)

        # train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
        # train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True, num_workers=2)

        return train_loader
    
    @pl.data_loader
    def val_dataloader(self):
      transform = transforms.Compose([
                                transforms.Resize(256),
                                transforms.CenterCrop(224),
                                transforms.ToTensor(),
                                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                                ])
                              
      val_set = torchvision.datasets.ImageFolder(os.path.join(self.data_dir, 'val'), transform)
      val_loader = torch.utils.data.DataLoader(val_set, batch_size=32, shuffle=True, num_workers=4)

      # val_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
      # val_loader = torch.utils.data.DataLoader(val_set, batch_size=4, shuffle=False, num_workers=2)

      return val_loader

    @pl.data_loader
    def test_dataloader(self):
      # The test set and validation set are the same so that I am able to check that the validation and test accuracy are equal.
      transform = transforms.Compose([
                              transforms.Resize(256),
                              transforms.CenterCrop(224),
                              transforms.ToTensor(),
                              transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                              ])
                            
      val_set = torchvision.datasets.ImageFolder(os.path.join(self.data_dir, 'val'), transform)
      val_loader = torch.utils.data.DataLoader(val_set, batch_size=4, shuffle=True, num_workers=4)

      return val_loader

I've also put my code into a Colab, if you would like to run it: https://colab.research.google.com/drive/1JUQctrkFKMJSPU2u5061534HzF27WFRE (the relevant part is in the first half of the notebook, up to 'For checking purposes'; the rest are messy and irrelevant, apologies)

Am I doing something wrong here, causing the accuracy returned to be different when I load the checkpoints separately instead of doing .test() immediately after .fit()?

[awaelchli]

Change
model_test.load_from_checkpoint('checkpoints/try_ckpt_epoch_1.ckpt')
to
model_test = CoolSystem.load_from_checkpoint('checkpoints/try_ckpt_epoch_1.ckpt')
and it will work (99.9% sure).

Note that load function is a class method. It will load the exact same hyperparameters from the checkpoint that you used for training.

[polars05]

Thanks, @awaelchli! Amended my code based on your suggestion and noted a strange behavior: during .fit(), the val accuracy is 0.8620 but upon loading the chekcpoints and doing .test(model_test), the test accuracy is 0.8813. This is strange to me as I'm using the same dataset for both val and test as a check.

[awaelchli]

Your notebook is quite big. It would be good if you could try to create a minimal example where this behavior can be observed, then it is easier for us to help. The difference in accuracy you report is quite small. Could it simply be due to randomization of the data? To make sure this doesn't happen, make the code deterministic: https://pytorch.org/docs/stable/notes/randomness.html

[polars05]

I'm loading my checkpoints for inference as per what @awaelchli suggested above:

from argparse import Namespace

args = {
    'num_classes': 2,
    'epochs': 2,
    'data_dir': "/content/hymenoptera_data",
}

hyperparams_test = Namespace(**args)

model_test = CoolSystem.load_from_checkpoint('checkpoints/try_ckpt_epoch_1.ckpt')
trainer = pl.Trainer(weights_summary=None)
trainer.test(model_test)

However, I am now getting an error:

AttributeError: 'dict' object has no attribute 'data_dir'

Even with model_test = CoolSystem(hyperparams_test).load_from_checkpoint('checkpoints/try_ckpt_epoch_1.ckpt'), PyTorch Lightning is still complaining that 'dict' object has no attribute 'data_dir'

Am I doing something wrong here?

[williamFalcon]

https://pytorch-lightning.readthedocs.io/en/0.7.1/weights_loading.html

[polars05]

I think the example shown in Checkpoint Loading only works if the user is loading the checkpoints immediately after training (i.e. the instance of MyLightningModule has already been instantiated with the necessary hparams).

However, I'm loading the checkpoints separately (i.e. not immediately after training), and I had to change __init__ within the module from

class CoolSystem(pl.LightningModule):
    def __init__(self, hparams):
        super(CoolSystem, self).__init__()

        self.hparams = hparams
        self.data_dir = self.hparams.data_dir 
        self.num_classes = self.hparams.num_classes

to

class CoolSystem(pl.LightningModule):
    def __init__(self, hparams):
        super(CoolSystem, self).__init__()

        self.hparams = hparams
        self.data_dir = self.hparams['data_dir']
        self.num_classes = self.hparams['num_classes']

in order to get the hyperparams to be loaded properly.

[williamFalcon]

your checkpoint was saved with a dictionary. means you likely gave the model a dictionary in init and not a namespace

[teichert]

@williamFalcon Could it be that this line is actually failing to convert the dictionary built by lightning back to a namespace. In particular, I believe that is happening to me because my checkpoint has no value for "hparams_type" which means that _convert_loaded_hparams gets a None as the second argument and returns the dictionary. In other words, the hparams in my checkpoint were indeed a namespace, but when trying to load the checkpoint, lightning builds a dictionary and it looks like the mechanism for changing it back to a namespace isn't working for me.

Possible work-arounds:

• Define a class variable hparams_type = 'hparams_name' tricks the conversion to return an AttributeDict (which is like a Namespace) [this isn't a reasonable work around, but I'm including it in case it helps to clarify what is going on]
• Check the type of hparams in __init__ and explicitly converting to a Namespace if necessary:

        if isinstance(hparams, dict):
            hparams = argparse.Namespace(**hparams)

(fwiw: I may be doing other things wrong, because my __init__ doesn't get any of the other arguments during reconstruction that it is supposed to get.) :)

[JanRuettinger]

I think I am facing a similiar issue. When I test a model which I load from a checkpoint file the results are quite bad.

TransferLearningModel.py

class TransferLearningModel(LightningModule):
    
    def __init__(self, config):
        super().__init__()
        
        self.config = config
        # save config dict when saving model
        self.save_hyperparameters(config)
       ....
        self.hparams.learning_rate = self.config["optimizer_initial_lr"]

        ......

training_testing.py

config = {
    "model_name": "resnext101_32x8d", 
    "amp_optimization_level": "O1", 
    ...
}

model = TransferLearningModel(config)

trainer = Trainer(...)
trainer.fit(model)

# testing
trainer.test(ckpt_path = filepath)

Any idea what I might be doing from?