LR Finder doesn't restore original model weights?

[AmarSaini]

Hey! I love this repo, thanks for making it 💯

Everything works well except for one thing, after some digging around/experimenting, here's what I've found:

Below are some figures for the training loss and training accuracy (on MNIST, using a resnet18).

Problem:

1. Using LRFinder on a model, and then training with it afterwards appears to hurt the models learning (see pink curve below).

Solution:

2. Using LRFinder on a model, and manually restoring the weights, appears to train the model optimally. (see green curve below).
3. Using LRFinder on a clone of the model, and then using the original model for training, appears to train the model optimally. (see green curve below).

Regarding the figure/graphs below, both models used the same hyperparameters.

An in-code example of option 1) would be similar to what was given in the README.md:

from torch_lr_finder import LRFinder

model = ...
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-7, weight_decay=1e-2)
lr_finder = LRFinder(model, optimizer, criterion, device="cuda")
lr_finder.range_test(trainloader, end_lr=100, num_iter=100)
lr_finder.plot()

// Then use "model" for training

An in-code example of option 3) would be:

from torch_lr_finder import LRFinder

model = ...
temp_model = *create model with same architecture*
// copy weights over
temp_model.load_state_dict(model.state_dict)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-7, weight_decay=1e-2)
// use temp model in lr_finder
lr_finder = LRFinder(temp_model, optimizer, criterion, device="cuda")
lr_finder.range_test(trainloader, end_lr=100, num_iter=100)
lr_finder.plot()

[NaleRaphael]

Hi @AmarSaini , thanks for these detailed report!

However, it seems like lr_finder.reset() wasn't called after lr_finder.range_test() finished. So that model and optimizer were not restored to their initial state. And that's probably the reason why the model performs bad after it.

Could you try to call lr_finder.reset() after the process of learning rate finding is finished, and check whether the model works normally or not?

If it still not work, please feel free to let us know!

[AmarSaini]

Wow, totally didn't see that haha. I did go through some of the source code and saw some code regarding restoring original model/opt parameters. I thought this was automatic (as in FastAI), my fault for not reading the complete README.md. Will try it again later tonight and follow up!

[AmarSaini]

Also, forgot to ask. In fastai we can run the lr_find after multiple training sessions.

With your API, we have to supply the optimizer ourselves. Just want to confirm that this means we should make a new (temp) optimizer, and copy the state from the current optimizer (used in training) over to the the new (temp) optimizer to use in LR_Find. (So lr_find is checking the loss with respect to different learning rates with the same optimizer state being used in training). This would be the correct thing to do right?

Or do you recommend using a fresh optimizer for LR_Find, even if you are planning to continue training?

Thanks :)

[NaleRaphael]

That's totally fine!
In my opinion, providing an API like lr_finder.reset() makes user able to do some post-processing (e.g. dump the weight for experimental test or something else) before the model and optimizer is restored.
As for the implementation in fastai, it's actually a great design for user to build a training pipeline.

For the question about using LRFinder in multiple training sessions, it doesn't need to make a new optimizer. But we have to create a new LRFinder in each new training session due to the implementation shown as below:

pytorch-lr-finder/torch_lr_finder/lr_finder.py

Lines 52 to 57 in 4de54de

	# Save the original state of the model and optimizer so they can be restored if
	# needed
	self.model_device = next(self.model.parameters()).device
	self.state_cacher = StateCacher(memory_cache, cache_dir=cache_dir)
	self.state_cacher.store('model', self.model.state_dict())
	self.state_cacher.store('optimizer', self.optimizer.state_dict())

Initial state of model and optimizer is only able to be cached when a LRFinder is instantiated. That's why we have to create a new LRFinder each time if we need to do learning rate finding in multiple training sessions (e.g. find learning rate after every 4 epochs of training).

In a nutshell, we should do like this in multiple training sessions:

# --- session 01 ---
lr_finder = LRFinder(...)
lr_finder.range_test(...)
lr_finder.reset()  # reset model and optimizer to their own initial state

# update learning rate of optimizer
best_lr = find_best_lr(lr_finder.history)  # a custom method to find best learning rate from history
for param_group in optimizer.param_groups:
    param_group['lr'] = best_lr

# start training model...


# --- session 02 ---
lr_finder = LRFinder(...)
lr_finder.range_test(...)
lr_finder.reset()

# update learning rate of optimizer
best_lr = find_best_lr(lr_finder.history)
for param_group in optimizer.param_groups:
    param_group['lr'] = best_lr

# start training model...

Hope these information help! :)

[AmarSaini]

Thanks for the detailed response :)

Regarding the multiple training sessions to continue training, is there functionality in this API to automatically reset an optimizer's hyperparameters?

(Before running session 2's LR Finder: I have to reset the learning rate to 1e-7. Can't just make a brand new optimizer, because we may want to maintain the optimizer state).

This is the correct thing to do right? (We have to make sure we have a min learning rate for the lr_range(), which by default uses what's in the optimizer)

E.g.

for param_group in optimizer.param_groups:
    param_group['lr'] = 1e-7
    // Also reset other optimizer hyperparameters to their defaults, e.g. momentum

[NaleRaphael]

Yes, you are right! We have to reset learning rate of the optimizer before running lr_finder.range_test() in the second session. Otherwise, lr_finder.range_test() will start the scanning from the latest learning rate used in previous session. (like the implementation in fastai.callbacks.lr_finder, learning rate will be reset to the initial value)

But I'm sorry about that there is currently no API to reset the hyperparameters of an optimizer automatically. Because that there are kinds of different implementation of optimizer, and we cannot determine what hyperparameters have to be reset. So this operation should be done by users.

However, it seems to be an interesting challenge to do so. As far as I know, fastai has implemented a wrapper for optimizer (see also: OptimWrapper).
I'll keep this in mind and maybe send a new PR for this if there is further breakthrough of design in torch.optim.

I'm glad to discuss with you! :)

[AmarSaini]

Thanks for all the info :)

And no worries, it's not too bad to reset the optimizer. As for it being trouble since different optimizers have different hyperparameters, here is how I would get around it:

optim_for_lr_find = torch.optim.*any optimizer with any settings*

optim_for_training = *some optim being used for training*

temp_dict = {}
temp_dict['state'] = optim_for_training.state_dict['state']
temp_dict['param_groups'] = optim_for_lr_find['param_groups']

optim_for_lr_find.load_state_dict(temp_dict)

// Run your LR Find with optim_for_lr_find

After plotting lr find, we can find a new optimal lr and put it in our current optimizer being used for training, without changing it's state:

optim_with_new_lr = torch.optim.*optimizer with optimal lr*

optim_for_training = *some optim being used for training*

temp_dict = {}
temp_dict['state'] = optim_for_training.state_dict['state']
temp_dict['param_groups'] = optim_with_new_lr['param_groups']

optim_for_training.load_state_dict(temp_dict)

// Train with optim_for_training

Note: an optimizer's state dict has a state and param_groups. Making the dictionaries manually lets you control which optimizer to take the state from, and which optimizer to take the (default) hyperparameters from (which we need a low learning rate for LR Find).

Thanks again for all your help!

[NaleRaphael]

That's a good idea!

And I've run some simple snippet to check the content of optimizer.param_groups:

# for `optim.SGD`
>>> optimizer.param_groups[0].keys()
dict_keys(['lr', 'momentum', 'dampening', 'weight_decay', 'nesterov', 'params'])

# for `optim.Adam`
>>> optimizer.param_groups[0].keys()
dict_keys(['params', 'lr', 'betas', 'eps', 'weight_decay', 'amsgrad', 'initial_lr'])

Since all items in optimizer.param_groups[0] (except optimizer.param_groups[0]['params']) are the hyperparameters of an optimizer, it seems that it's possible to make a generalized solution to reset them.

But I have to take more time to dive into the source code of pytorch to figure it out.
Thanks for your help! ;)

[AmarSaini]

Oh nice!

Let me know if you're looking for any help, I'd be more than happy to contribute on any part of this lovely repo! :)

[NaleRaphael]

Sure! I'll update this thread once there is any breakthrough!

[davidtvs]

Is there a need to reset other parameters other than the learning rate?

Seems to me that you would always want to keep everything else as-is and just lower the learning rate. If that's the case then we can add an initial_lr parameter to range_test. See the code example below changed from what @NaleRaphael posted before:

# --- session 01 ---
optimizer = torch.optim... # Some optimizer
lr_finder = LRFinder(model, optimizer, ...)
lr_finder.range_test(initial_lr=None, ...) # Set to None defaults to the lr in optimizer
lr_finder.reset()  # reset model and optimizer to their own initial state

# update learning rate of optimizer
best_lr = find_best_lr(lr_finder.history)  # a custom method to find best learning rate from history
for param_group in optimizer.param_groups:
    param_group['lr'] = best_lr

# start training model...


# --- session 02 ---
lr_finder = LRFinder(model, optimizer, ...) # optimizer was used for training so it carries the state from training
lr_finder.range_test(initial_lr=1e-7, ...) # change the initial learning rate to something else
lr_finder.reset() # resets the optimizer to the state and parameters after training

# update learning rate of optimizer
best_lr = find_best_lr(lr_finder.history)
for param_group in optimizer.param_groups:
    param_group['lr'] = best_lr

[AmarSaini]

Hi! I'm not too sure how important it is if we preserve the other hyper-parameters in the optimizer. Currently, I'm using the OneCycle LR Scheduler with the Adam Optimizer.

Before training this is what my optimizer looks like for using LR Finder (no lr scheduler), using lr=1e-7:

After find a good learning rate, lr=1e-3 training this is what my optimizer looks after wrapped with the OneCycle LR Scheduler:

After training session 1, this is what my optimizer looks like:

Before training session 2, let's say we want to find another learning rate.
Will a wrapped LR Scheduler effect the LR Finder?
Simply reseting the lr back to lr=1e-7 would yield this for the optimizer:

Running Find LR with this resetted lr=1e-7 optimizer wrapped with a LR Scheduler worked fine:

Then we can set the lr to something like lr=1e-3

Then finally, re-wrap with the OneCycle LR Scheduler to get:

So it looks like this would work nicely. My current solution was just to make a temp optimizer to get default param groups, and manually copy the state over to it.

[NaleRaphael]

@davidtvs :

Is there a need to reset other parameters other than the learning rate?

If there is an optimizer of which the hyperparameters will change according to iteration, then we might need to reset it. But I am not sure whether there is a kind of this optimizer up to now.

Seems to me that you would always want to keep everything else as-is and just lower the learning rate. If that's the case then we can add an initial_lr parameter to range_test.

This sounds great!

@AmarSaini :

Will a wrapped LR Scheduler (optimizer?) effect the LR Finder?

As far as I know, lr of an optimizer which is wrapped by torch.optim.lr_scheduler.* will change in only 2 stage:

1. when lr_scheduler.step() is called (see also torch.optim.lr_scheduler#L88-89)
2. when a lr_scheduler is just initialized (see also torch.optim.lr_scheduler#L43)

With the point 2 mentioned above, we should take care that the beginning lr used by lr_finder will be affected if an optimizer has been wrapped by a lr_scheduler before running lr_finder.range_test().
For this reason, I would suggest performing lr_finder.range_test() before wrapping a optimizer by any lr_scheduler to avoid unexpected result.

[AmarSaini]

@davidtvs @NaleRaphael

For this reason, I would suggest performing lr_finder.range_test() before wrapping a optimizer by any lr_scheduler to avoid unexpected result.

Would this mean the end-user should be responsible for un-wrapping an optimizer before calling lr_finder.range_test() a second time before another training session?

I've been playing around with un-wrapping optimizers, stripping away their lr_schedulers (by making new state_dicts with the state from the main optimizer being used for training, and the param_groups of an optimizer freshly made without it being wrapped in a lr_scheduler), and so far I found that just making a new optimizer for a new training session works fine (only on MNIST though haha)

[davidtvs]

@AmarSaini Ya, currently, the end-user should take care of it.
We could change that in the future though; check if the optimizer has a scheduler wrapped around it and do something like what you described. I can look into it, or review a PR for this