In a transfer learning setting, I want to freeze the body and only train the head for 2 epochs. Then I want to unfreeze the whole network and use the Learning Rate finder, before continue training again.
What I want to do is similar to FastAI’s
To do the same with PyTorch Lightning, I tried the following:
Trainer(max_epochs=2, min_epochs=0, auto_lr_find=True) trainer.fit(model, data_module) # FastAI: learn.fit_one_cycle(2) trainer.max_epochs = 5 # model.unfreeze() # allow the whole body to be trained # trainer.tune(model) # LR finder trainer.fit(model, data_module) # FastAI: learn.fit_one_cycle(3)
Unfortunately, this would invoke the training op
epoch 1 twice.
Any ideas of how to approach what I want to do?
Edit 1: Link to Colab demonstrating 2x epoch 1: https://colab.research.google.com/drive/15yYTPSrv1e4yUDKNGtSp8A3WqLmBzATF?usp=sharing