OK, you know how time consuming the bs discovery is, you got to run lr_find/fit and then if bs is too big do it again, then the next stage, and repeat it all together again and again, then the next stage, and meanwhile waiting, and waiting and waiting…
Then recently I discovered that all the required GPU memory is setup during the first batch of the first epoch see tutorial. So I thought why do I waste hours waiting for things to finish. I also find the waiting makes me really tired physically. Surely, I can truncate my dataset while tuning things up, but even that is not very efficient if say the batch size progressively gets smaller and image size gets larger and it’s all error-prone and you have to manually change things and to remember to put them back to normal… too much hassle.
So Sylvain and I came up with a quickie method, all you do is add the following code somewhere in your notebook before your first learn is created:
# this is for bs tune up - it will run through fit calls really fast, doing the minimum required to setup gpu mem requirements
from fastai.callback import Callback
class StopAfterNBatches(Callback):
"Stop training after n batches of the first epoch."
def __init__(self, n_batches:int=2):
self.stop,self.n_batches = False,n_batches-1 # iteration starts from 0
def on_batch_end(self, iteration, **kwargs):
if iteration == self.n_batches: return {'stop_epoch': True, 'stop_training': True}
This version works in v1.0.47 and later. It’s part of the library now.
https://docs.fast.ai/callbacks.misc.html#StopAfterNBatches
So load it with:
from fastai.callbacks.misc import *
And then to use it with your learn object:
from fastai.callbacks.misc import *
learn.callbacks.append(StopAfterNBatches(n_batches=2))
or even better you can set it globally for all learners in that notebook, without needing to add/remove any code:
defaults.extra_callbacks = [StopAfterNBatches(n_batches=2)]
And here is how to make it easy to turn it on and off globally in your notebook:
# True turns the speedup on, False return to normal behavior
tune = True
if tune:
defaults.extra_callbacks = [StopAfterNBatches(n_batches=2))]
else:
defaults.extra_callbacks = None
This callback, once set, will impact all the subsequent lr_find and fit* calls. They will all finish within seconds and fail as normal if there isn’t enough RAM, which is when you reduce your bs and other params.
Remember this is for quick batch size (and other param) tuning, ignore the fact that the training data is completely premature and you will have all progress bars red as they all get interrupted.
When done tuning, flip the tune var to False and everything is back to normal. If later you need to re-tune since say you made major changes, flip tune on again, run the nb, adjust params, and turn it off again.
I’ve only tried it in a few vision notebooks and the memory consumption is almost identical as the full normal run.
So please give it a test run, I haven’t tried it with any text training for example.
Some learners will certainly have minor issues. For example, unet_learner currently has a pretty significant memory consumption fluctuation on each run (even with fixed seed), so beware of that - unet learner is not stable memory-wise at the moment.
Feedback and improvement suggestions are welcome.
If it looks like a solid tech we will integrated it into fastai. The defaults.callbacks is already in, I am referring to adding the special callback so that you don’t have to copy-n-paste it.
And the next stop is batch-size and other hyper-param auto-discoverer, which some of you have already started working on.
If you’d like to discuss the methods of choosing a better bs, please see this thread.
