How to ignore certain "parameters" with model checkpointing?

[jmerkow]

I have some data that I store on my LightningModule during validation. I want to prevent this from being saved by the model checkpoint. They are not actually parameters and do not affect the state at all. I want to maintain other parts of the state, I don't want to use weights only.

Is it possible to do this?

[ananthsub]

Out of curiosity, what are the states? You could try removing these states within the LightningModule's on_save_checkpoint hook

[jmerkow]

They are not actually states. I have an autoencoder, and I am saving the best and worst images. To make things faster, they are stored as tensors, I think trainer is picking them up and trying to save them. Essentially, I am getting RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! and trying to track down what on the wrong device. I am not 100% sure whats causing the issue, but I think it these images.
Is there a way to see specifically which device each thing being saved is on? i.e.

class ModelCheckpointDBG(ModelCheckpoint):
    def save_checkpoint(self, trainer, unused=None):
        try:
            super().save_checkpoint(trainer, unused=unused)
        except RuntimeError:
            ## What goes here??
            raise

Either way I will need to ignore them

[jmerkow]

Wait a second, it appears that this isn't part of the save.

[tchaton]

Hey @jmerkow,

The checkpoint is generated from the dump_checkpoint function of the CheckpointConnector. One of the latest hooks called is lightning_module.on_save_checkpoint(checkpoint) here: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/connectors/checkpoint_connector.py#L386

This is actually done in-place.

Therefore, you could do the following.

class MyModel(LightningModule):

    def on_save_checkpoint(self, checkpoint):
        # pop the keys you are not interested by
        checkpoint

[jmerkow]

I ended up realizing I was having a different problem, but this answers my question, ty