Evaluating on multiple tasks in continual learning setup.

[prateeky2806]

Hi Everyone, this is my first project in PyTorch lightning and I am working on a problem related to continual learning. During training I have multiple Tasks (Datasets A, B, C, D) which are used sequentially in the training. Now, whenever I am switching from task x to task y, I want to evaluate the performance on the validation on all the previous and future tasks.

My Datamodule looks something like this, where I have a method update_task(task_id) which updates the dataloaders to the correct task and also resets the train and val dataloaders whenever it is called.

class ContinualDataset(LightningDataModule):
    def __init__(self):
        .....

    def setup(self, stage: Optional[str] = None):
      
        datasets = [...]

        self.loaders = [
            (DataLoader(train_data, batch_size=self.hparams.batch_size, shuffle=True),
            DataLoader(val_data, batch_size=self.hparams.test_batch_size, shuffle=False),)
            for train_data, val_data in datasets
        ]

    def update_task(self, i):
        self.current_task = i
        self.trainer.reset_train_dataloader()
        self.trainer.reset_val_dataloader()
        
    def train_dataloader(self):
        return self.loader[self.current_task][0]

    def val_dataloader(self):
        return self.loader[self.current_task][1]

I am using a logic like this to perform the training. I am using on_train_epoch_start to switch the dataloaders every k epochs.

class LitModel(LightningModule):
    def on_train_epoch_start(self):
        if self.trainer.current_epoch == 0 or (self.trainer.current_epoch + 1) % self.trainer.reload_dataloaders_every_n_epochs == 0:
            # update task and model parameters            
            self.trainer.datamodule.update_task(self.task_idx)
            self.update_model_parama(self.task_idx)

Now after every k epoch, Now from epoch 1-k I am performing validation on task 1 but at the kth epoch I also want to validate on all the other tasks as well (A,B,C,D). In pure pytorch it would look something like this.

def validate_on_all_tasks():
        acc_object = Acc()
        Acc_dict = {}
        for idx in range(num_tasks):
            # also reloads the dataloaders
            datamodule.update_task(idx)
            for batch in datamodule.val_loader:
                x, y = batch
                logits = self.model(x)
                preds = torch.argmax(logits, dim=-1)
                acc_obj.update(preds, y)
            Acc_dict[f'Acc:{idx}'] = acc_obj.compute()
            acc_object.reset()
        return Acc_dict

for epoch in range(num_epochs):
    train(...)
    validate_on_current_task(...)
    if epoch % k =0:
        performance_on_all_tasks = validate_on_all_tasks()
    

I am not sure how to go about this in pytorch lightning, specifically how to get manually iterate over dataloader in lightning as currently I just work wby defining validation and training steps which takes care of iterating over the dataloader internally. I tried to do it myself manually but then I get errors regarding tensors on being on right devices. Any help would be highly appreciated! Thanks in Advance.

[rohitgr7]

for such a case, I'd suggest extracting the loaders from your datamodule and passing each dataloader from 1 to k explicitly.

for i in range(1, k):
     val_dl = datamodule.loaders[i]
     res = trainer.validate(model, dataloaders=val_dl)