How to change EvaluationLoop for time dependent data

[Udneowi]

Hi!

I am trying to rewrite some old code I wrote in vanilla PyTorch to PL, but I am having some issues with the fact that I am dealing with time dependent data.

In short the network is predicting the next token in a sequence, where each sequence is independent and of varying lengths. During validation I therefor have to warm up the hidden state for each sequence and then begin predictions. Old pseudo code was something like the following:

#validation_loop
for sequence_id in sequences.keys():
    hidden = init_hidden_state() 

    # Running model through the history of the sequence to warmup the hidden state
    for step, (input, target) in enumerate(dataloader_warmup(data[sequence_id])):
        _, hidden, = model(input, hidden)
    # Doing predictions on future unseen data with the warmed up hidden state
    for step, (input, target) in enumerate(dataloader_prediction(data[sequence_id])):
        out, hidden, = model(input, hidden)
        # Compute loss and accuracy of validation step

However in PyTorch Lightning it is not straight forward to copy this "double loop" with warmup format. I have tried creating my own EvaluationLoop after reading the page on Loops where it seemed simple but with no success. My hacky attempt was the following:

class LSTM_validation_loop(EvaluationLoop):

    def run(self):
        count = 0
        # Looping over sequence ID
        for pep in self.dataloaders[0]:
            # Initialize hidden state
            self.trainer.lightning_module.net.hidden = self.trainer.lightning_module.net.init_hidden(bs = 1)
            # Loop over warmup data to warmup the hidden state
            for (inp, _, _, _ )in self.dataloaders[0].dataset.val_warm_dat(int(pep), "warm"):
                _, self.trainer.lightning_module.net.hidden, _ = self.trainer.lightning_module.net(inp)
            # Use the warmed hidden state to do predictions
            for batch in self.dataloaders[0].dataset.val_warm_dat(int(pep), "val"):
                loss = self.trainer.lightning_module.validation_step(batch, count)
                count += 1

which does "work" in the sense that it computes all the relevant stuff, but it messes to much with the underlying functions and I have no clue where to store the output.

Does anyone have a lead on how to solve this? :)

Any ideas is much appreciated!