How to include validation dataset for a final training epoch

[hogru]

Hello, can you please put me on the right trajectory (link, method, ...) for the following: after training I would like to do a final (number of) training epochs including the validation dataset. I read about concatenating datasets and combining dataloaders, but how could I use those for an additional set of training epochs only?

[rohitgr7]

hey @hogru !

can you explain more about what are you trying to do here? step-by-step workflow on how it should look like?

[hogru]

Hi @rohitgr7,

the idea/practice is the following: I train a model "as usual": training loop, after each epoch I do a validation, I save the model if it's better than the previous best, do early stopping etc. The result is my current best model trained on the training set. Pretty standard until this point I guess.

What comes next is the core of my question: I take this best model and train it for one (or a very small number of) epochs on the combined training and validation set. This way, the model sometimes improves a bit (depending on split sizes, moon phases, etc.)

In vanilla pytorch code this would look something like:

# Model ist trained and saved in "best_model"; load it for a final epoch
model.load_state_dict(best_model.state_dict())

# Combine the training and validation dataloaders, e.g. with chain from Iterables
total_loader = chain(val_loader, train_loader)

# Train for 1 final epoch loading data from total_loader and save model 

Does that help?

Best,
Stephan

[rohitgr7]

maybe you can configure your datamodule with a flag

class LitDM(LightningDataModule):
    def __init__(self, combine=False):
        ...

    def train_dataloader(self):
         if combine:
             return chain(train_dataloader, val_dataloader)

         return train_dataloader

first training:

trainer = Trainer(...)
model = LitModel()
dm = LitDM(combine=False)
trainer.fit(model, datamodule=dm)

2nd training:

trainer = Trainer(...)
model = LitModel.load_from_checkpoint('best_model.ckpt')
dm = LitDM(combine=True)
trainer.fit(model, datamodule=dm)

[crazyn2]

It didn't work. The train_dataloader(self) didn't replace the pytorch lightning tran_dataloader of trainer.