How to checkpoint model predictions?

1

Hello,

I would like to save some of my model outputs to file. This will be done only at the test step:

    def test_step(self, batch, batch_idx):
        x1, x2 = batch["x1"], batch["x2"]
        predicts = self(x1, x2) # which call model forward in eval mode
        save(predicts, predicts_path) # method to save predicts to file

What would be an efficient approach to concatenate in a file each batch predicts of the model (which has the shape [batch_size, hidden_size])?

2

The easiest approach is to use torch.save and torch.load that will handle the torch.Tensor serialization and deserialization respectively.

With this in mind, here are the functions. The first one appends the predictions of a model and the second one reads them in the form of a list.

def save_data(data, path):
    if isinstance(data, torch.Tensor):
        data = [data]
    
    with open(path, 'ab') as f:
        for d in data:
            torch.save(d, f)

def load_data(path):
    data = []
    with open(path, 'rb') as f:
        try:
            while True:
                data.append(torch.load(f))
        except EOFError:
            pass    

    return data