Writing predictions for training and validation dataloaders on fit end

[adamklie]

Is there a simple way to take advantage of current Lightning features to write predictions on the entire training and validation sets to separate files on fit end? Been trying to take advantage of some combination of the on_fit_end and BasePredictionWriter (which I am already using for writing using predict functionality). Worst case I will just write a separate script but wanted to see if there was something built in I could take advantage of.

Thanks!

[rohitgr7]

there is no built-in component for it and baseprediction writer callback is the best one you can configure and use.

[rohitgr7]

if you are doing epoch based training then you can do something like:

class MyPredicitonWriter(Callback):
    def __init__(self):
        self.train_preds = []
        self.val_preds = []
 
    def on_train_batch_end(self, trainer, pl_module, outputs, *args, **kwargs):
        if trainer.current_epoch == trainer.max_epochs - 1:
            self.train_preds.append(outputs)

    def on_train_epoch_end(self, trainer, pl_module):
        if trainer.current_epoch == trainer.max_epochs - 1:
            # save train_preds to file

    def on_validation_batch_end(self, trainer, pl_module, outputs, *args, **kwargs)
         if trainer.current_epoch == trainer.max_epochs - 1:
             self.val_preds.append(outputs)
   
    def on_validation_epoch_end(self, trainer, pl_module):
         if trainer.current_epoch == trainer.max_epochs - 1:
             # save val_preds to file


trainer = Trainer(..., callbacks=[MyPredicitonWriter()])

[adamklie]

Ok got it. Does that mean that the training_step and val_step functions in my pl_module need to return what I expect to output to the file. Where would the outputs arguments in the above functions be coming from?

[rohitgr7]

yes.. the stuff that you return from the training/validation_step hooks.

[adamklie]

Sorry for the naivete here, but how would that work for the training step hook? Would I have to override the default FitLoop to extract the loss from the training step return? Just trying to make sure I'm not missing anything.

[rohitgr7]

Sorry, didn't get you. Can you explain with some pseudo code?

[adamklie]

Sure. As far as I understand, the default loops called when running fit in pl expect the training step to return just the loss. Something like:

# FitLoop
for epoch in range(max_epochs):

    # TrainingEpochLoop
    for batch_idx, batch in enumerate(train_dataloader):

        # TrainingBatchLoop
        for split_batch in tbptt_split(batch):

            # OptimizerLoop
            for optimizer_idx, opt in enumerate(optimizers):

                loss = lightning_module.training_step(batch, batch_idx, optimizer_idx)

But in order to write predictions for the training data using something similar to the MyPredictionWriter you outlined above, I would need the lightning_module's training_step to return not only the loss, but also the predictions. So the training_step would return the loss and predictions in some kind of data structure

def training_step(self, batch, batch_idx):
        # Compute predictions and loss
        return loss, predictions

So I would have thought I would need override this loop in order to handle the fact that the training_step hook returns more than just the loss. Something like:

# FitLoop
for epoch in range(max_epochs):

    # TrainingEpochLoop
    for batch_idx, batch in enumerate(train_dataloader):

        # TrainingBatchLoop
        for split_batch in tbptt_split(batch):

            # OptimizerLoop
            for optimizer_idx, opt in enumerate(optimizers):

                loss = lightning_module.training_step(batch, batch_idx, optimizer_idx)[0]  # Grab the first index of the returned values

Does that make sense?

[ndahiya3]

@rohitgr7 Yeah, return predictions say for say only 10 batches per epoch to save on hard drive. I can't save all the predictions on either GPU or CPU as dataset is large.

[rohitgr7]

@ndahiya3

you can just add conditions based on batch_idx.

def training_step(self, batch, batch_idx):
    preds = ...
    if batch_idx < 10:
        return {'loss': loss, 'preds': preds}
    return {'loss': loss'}

def training_epoch_end(self, outputs):
    for out in outputs:
        if 'pred' in out:
            save(out['pred'])

[ndahiya3]

@rohitgr7 Thanks. I think I'll do something like you mentioned. I am getting following warning: ~python3.9/site-packages/pytorch_lightning/loops/optimization/closure.py:35: LightningDeprecationWarning: One of the returned values {'preds'} has a grad_fn. We will detach it automatically but this behaviour will change in v1.6. Please detach it manually: return {'loss': ..., 'something': something.detach()}

Should prediction be detached in the training_step() ? Loss.backward() is yet to be run at this stage, right? If I detach prediction at this stage will it not interfere with the backward pass running correctly? Detached prediction will disconnect the loss from the rest of the network, I believe.

[semaphore-egg]

@ndahiya3 Unless you have some advanced purpose, let us detach every tensor, except the loss, to be returned from training_step.
detached_t=source_t.detach() means creating a new tensor detached_t which shares inner data with source_t. I.e., shallow copy.

The source tensor source_t is still in the computation graph and its gradients are recorded once further differentiable operations are applied.

The new tensor detached_t is not in the computational graph. Operations on detached_t do not change the computational graph, except that directly modifying data of detach_t means modifying data of source_t.

A safer way is to make a deep copy detached_t=source_t.detach().clone(). In this way, detached_t and source_t are totally unrelated.

[rohitgr7]

@ndahiya3

If I detach prediction at this stage will it not interfere with the backward pass running correctly?

it won't. You can safely detach it since loss is already connected to the graph.