<jupyter_start><jupyter_text>Simple Profiler> Callbacks which add a simple profiler to fastai. Inspired by PyTorch Lightning's SimpleProfiler. > Note: Simple Profiler is currently untested on distributed training.  EventsFastai callbacks do not have an event which is called directly before drawing a batch. Simple Profiler adds a new callback event called `before_draw`. With Simple Profiler imported, a callback can implement actions on the following events:- `after_create`: called after the `Learner` is created- `before_fit`: called before starting training or inference, ideal for initial setup.- `before_epoch`: called at the beginning of each epoch, useful for any behavior you need to reset at each epoch.- `before_train`: called at the beginning of the training part of an epoch.- `before_draw`: called at the beginning of each batch, just before drawing said batch. - `before_batch`: called at the beginning of each batch, just after drawing said batch. It can be used to do any setup necessary for the batch (like hyper-parameter scheduling) or to change the input/target before it goes in the model (change of the input with techniques like mixup for instance).- `after_pred`: called after computing the output of the model on the batch. It can be used to change that output before it's fed to the loss.- `after_loss`: called after the loss has been computed, but before the backward pass. It can be used to add any penalty to the loss (AR or TAR in RNN training for instance).- `before_backward`: called after the loss has been computed, but only in training mode (i.e. when the backward pass will be used)- `before_step`: called after the backward pass, but before the update of the parameters. It can be used to do any change to the gradients before said update (gradient clipping for instance).- `after_step`: called after the step and before the gradients are zeroed.- `after_batch`: called at the end of a batch, for any clean-up before the next one.- `after_train`: called at the end of the training phase of an epoch.- `before_validate`: called at the beginning of the validation phase of an epoch, useful for any setup needed specifically for validation.- `after_validate`: called at the end of the validation part of an epoch.- `after_epoch`: called at the end of an epoch, for any clean-up before the next one.- `after_fit`: called at the end of training, for final clean-up.  Implement before_draw To add `before_draw` as a callable event, first it needs to be added to both lists of fastai events.<jupyter_code>#exports
_inner_loop = "before_draw before_batch after_pred after_loss before_backward before_step after_step after_cancel_batch after_batch".split()
#exports
_events = L.split('after_create before_fit before_epoch before_train before_draw before_batch after_pred after_loss \
    before_backward before_step after_cancel_step after_step after_cancel_batch after_batch after_cancel_train \
    after_train before_validate after_cancel_validate after_validate after_cancel_epoch \
    after_epoch after_cancel_fit after_fit')

mk_class('event', **_events.map_dict(),
         doc="All possible events as attributes to get tab-completion and typo-proofing")<jupyter_output><empty_output><jupyter_text>Next, `Callback` needs to be modified to be aware of the new event.<jupyter_code>#exports
@patch
def __call__(self:Callback, event_name):
    "Call `self.{event_name}` if it's defined"
    _run = (event_name not in _inner_loop or (self.run_train and getattr(self, 'training', True)) or
            (self.run_valid and not getattr(self, 'training', False)))
    res = None
    if self.run and _run: res = getattr(self, event_name, noop)()
    if event_name=='after_fit': self.run=True #Reset self.run to True at each end of fit
    return res<jupyter_output><empty_output><jupyter_text>Then `Learner._call_one` needs to patch to be aware of the `before_draw`.<jupyter_code>#exports
@patch
def _call_one(self:Learner, event_name):
    if not hasattr(event, event_name): raise Exception(f'missing {event_name}')
    for cb in self.cbs.sorted('order'): cb(event_name)<jupyter_output><empty_output><jupyter_text>Finally, `Learner.all_batches` can be modified to call `before_draw` when iterating through a dataloader.<jupyter_code>#exports
@patch
def all_batches(self:Learner):
    self.n_iter = len(self.dl)
    if hasattr(self, 'simple_profiler'):
        self.it = iter(self.dl)
        for i in range(self.n_iter):
            self("before_draw")
            self.one_batch(i, next(self.it))
        del(self.it)
    else:
        for o in enumerate(self.dl): self.one_batch(*o)<jupyter_output><empty_output><jupyter_text>While testing hasn't shown any negative side effects of this approach, `all_batches` only uses the new batch drawing implementation if `SimpleProfilerCallback` is in the list of callbacks, and reverts back to the original method if not.<jupyter_code>#export
_loop = ['Start Fit', 'before_fit', 'Start Epoch Loop', 'before_epoch', 'Start Train', 'before_train',
         'Start Batch Loop', 'before_draw', 'before_batch', 'after_pred', 'after_loss', 'before_backward', 
         'before_step', 'after_step', 'after_cancel_batch', 'after_batch','End Batch Loop', 'End Train',
         'after_cancel_train', 'after_train', 'Start Valid', 'before_validate', 'Start Batch Loop',
         '**CBs same as train batch**', 'End Batch Loop', 'End Valid', 'after_cancel_validate',
         'after_validate', 'End Epoch Loop', 'after_cancel_epoch', 'after_epoch', 'End Fit',
         'after_cancel_fit', 'after_fit']
#exporti
@patch
def show_training_loop(self:Learner):
    indent = 0
    for s in _loop:
        if s.startswith('Start'): print(f'{" "*indent}{s}'); indent += 2
        elif s.startswith('End'): indent -= 2; print(f'{" "*indent}{s}')
        else: print(f'{" "*indent} - {s:15}:', self.ordered_cbs(s))<jupyter_output><empty_output><jupyter_text>Simple Profiler Callbacks -<jupyter_code>#exporti
_all = ['fit', 'epoch', 'train', 'validate']
_train = ['backward', 'step', 'zero_grad']
_multiple = ['draw', 'batch', 'pred', 'loss']
#export
class SimpleProfilerPostCallback(Callback):
    "Pair with `SimpleProfilerCallback` to profile training performance. Removes itself after training is over."
    order,remove_on_fetch = Recorder.order-1,True

    def before_fit(self):
        self.profiler = self.learn.simple_profiler

    def after_train(self):
        self.profiler._raw_time['train'].append(time.monotonic() - self.profiler._train_start)

    def after_validate(self):
        self.profiler._raw_time['validate'].append(time.monotonic() - self.profiler._validate_start)

    def after_pred(self):
        if self.training: self.profiler._raw_time['train_pred'].append(time.monotonic() - self.profiler._train_batch_start)
        else: self.profiler._raw_time['valid_pred'].append(time.monotonic() - self.profiler._valid_batch_start)
        
        if self.training: self.profiler._train_loss_start = time.monotonic()
        else: self.profiler._valid_loss_start = time.monotonic()

    def after_loss(self):
        if self.training: self.profiler._raw_time['train_loss'].append(time.monotonic() - self.profiler._train_loss_start)
        else: self.profiler._raw_time['valid_loss'].append(time.monotonic() - self.profiler._valid_loss_start)

    def after_step(self):
        self.profiler._raw_time['step'].append(time.monotonic() - self.profiler._step_start)
        self.profiler._zero_start = time.monotonic()

    def after_batch(self):
        if self.training: self.profiler._raw_time['train_batch'].append(time.monotonic() - self.profiler._train_draw_start)
        else: self.profiler._raw_time['valid_batch'].append(time.monotonic() - self.profiler._valid_draw_start)

    def after_epoch(self):
        self.profiler._raw_time['epoch'].append(time.monotonic() - self.profiler._epoch_start)

    def after_fit(self):
        self.profiler._raw_time['fit'].append(time.monotonic() - self.profiler._fit_start)
        self.profiler._generate_report()
        if not hasattr(self.learn, 'lr_finder'):
            self.profiler._display_report()
            self.learn.remove_cbs([SimpleProfilerCallback, SimpleProfilerPostCallback])
#export
class SimpleProfilerCallback(Callback):
    """
    Adds a simple profiler to the fastai `Learner`. Optionally showing formatted report or saving unformatted results as csv.
    
    Pair with SimpleProfilerPostCallback to profile training performance.

    Post fit, access report & results via `Learner.simple_profile_report` & `Learner.simple_profile_results`.
    """
    order,remove_on_fetch = TrainEvalCallback.order+1,True
    def __init__(self, show_report=True, plain=False, markdown=False, 
                 save_csv=False, csv_name='simple_profile.csv'): 
        store_attr()
        self.csv_name = Path(csv_name)

    def before_fit(self):
        self._raw_time = dict()
        for c in _all + _train:
            self._raw_time[c] = []
        for p in ['train', 'valid']:
            for c in _multiple:
                self._raw_time[f'{p}_{c}'] = []

        self._fit_start = time.monotonic()

    def before_epoch(self):
        self._epoch_start = time.monotonic()

    def before_train(self):
        self._train_start = time.monotonic()

    def before_validate(self):
        self._validate_start = time.monotonic()

    def before_draw(self):
        if self.training: self._train_draw_start = time.monotonic()
        else: self._valid_draw_start = time.monotonic()

    def before_batch(self):
        if self.training: self._raw_time['train_draw'].append(time.monotonic() - self._train_draw_start)
        else: self._raw_time['valid_draw'].append(time.monotonic() - self._valid_draw_start)

        if self.training: self._train_batch_start = time.monotonic()
        else: self._valid_batch_start = time.monotonic()

    def before_backward(self):
        self._backward_start = time.monotonic()

    def before_step(self):
        self._raw_time['backward'].append(time.monotonic() - self._backward_start)
        self._step_start = time.monotonic()
    
    def after_batch(self):
        if self.training: self._raw_time['zero_grad'].append(time.monotonic() - self._zero_start)

    def _generate_report(self):
        total_time = self._raw_time['fit'][0]
        self.report = pd.DataFrame(columns=['Phase', 'Action', 'Mean Duration', 'Duration Std Dev', 
                                               'Number of Calls', 'Total Time', 'Percent of Total'])

        for c in _all:
            if c == 'fit':
                self._append_to_df(['fit', c, 0, 0, 1, total_time, f'{self._calc_percent(total_time):.0%}'])
            else: self._append_to_df(self._create_row('fit', c, self._raw_time[c]))
        
        for c in _train: self._append_to_df(self._create_row('train', f'{c}', self._raw_time[c]))

        for p in ['train', 'valid']:
            for c in _multiple: self._append_to_df(self._create_row(p, f'{c}', self._raw_time[f'{p}_{c}']))

        self.report = self.report.sort_values(['Phase','Total Time'], ascending=[True,False])
        self.learn.simple_profile_results = self.report.copy()
        for c in ['Mean Duration', 'Duration Std Dev', 'Total Time']:
            self.report[c] = self.report[c].apply(self._scale)
        self.report['Phase'] = self.report['Phase'].where(~self.report['Phase'].duplicated(), '')

        self.learn.simple_profile_report = self.report

    def _display_report(self):   
        if self.show_report:
            if self.markdown: print(self.report.to_markdown(index=False))
            else:
                if in_notebook() and not self.plain:
                    with pd.option_context('display.max_rows', len(self.report.index)):
                        s = self.report.style.set_caption("Simple Profiler Results").hide_index()
                        display(s)
                else:
                    print('Simple Profiler Results')
                    print(self.report.to_string(index=False))

        if self.save_csv:
            self.path.parent.mkdir(parents=True, exist_ok=True)
            self.learn.simple_profile_results.to_csv(self.path/self.csv_name, index=False)

    def _append_to_df(self, row):
        self.report.loc[len(self.report.index)] = row

    def _calc_percent(self, time):
        return time / self._raw_time['fit'][0]

    def _create_row(self, phase, name, input):
        return [phase, name, np.mean(input), np.std(input), len(input), 
                np.sum(input), f'{self._calc_percent(np.sum(input)):.0%}']

    # modified from https://github.com/thomasbrandon/mish-cuda/blob/master/test/perftest.py
    def _scale(self, val, spec="#0.4G"):
        if val == 0: return '-'
        PREFIXES = np.array([c for c in u"yzafpnµm kMGTPEZY"])
        exp = np.int8(np.log10(np.abs(val)) // 3 * 3 * np.sign(val))
        val /= 10.**exp
        prefix = PREFIXES[exp//3 + len(PREFIXES)//2]
        return f"{val:{spec}}{prefix}s"<jupyter_output><empty_output><jupyter_text>Convenience Method<jupyter_code>#export
@patch
def profile(self:Learner, show_report=True, plain=False, markdown=False, 
            save_csv=False, csv_name='simple_profile.csv'):
    "Run Simple Profiler when training. Simple Profiler removes itself when finished."
    self.add_cbs([SimpleProfilerCallback(show_report, plain, markdown, save_csv, csv_name),
                  SimpleProfilerPostCallback()])
    return self<jupyter_output><empty_output><jupyter_text>Output The Simple Profiler report contains the following items divided in three Phases (Fit, Train, & Valid)Fit:- `fit`:  total time fitting the model takes.- `epoch`: duration of both training and validation epochs. Often epoch total time is the same amount of elapsed time as fit.- `train`: duration of each training epoch.- `validate`: duration of each validation epoch.Train:- `draw`: time spent waiting for a batch to be drawn. Measured from `before_draw` to `before_batch`. With default prefetching settings, ideally this is as close to instantly as possible. - `batch`: total duration of all batch steps sans drawing the batch. Measured from `before_batch` to `after_batch`.- `pred`: duration of the forward pass and any additional batch modifications. Measured from `before_batch` to `after_pred`.- `loss`: duration of caculating loss. Measured from `after_pred` to `after_loss`.- `backward`: duration of the backward pass. Measured from `before_backward` to `before_step`.- `step`: duration of the optimizer step. Measured from `before_step` to `after_step`.- `zero_grad`: duration of the zero_grad step. Measured from `after_step` to `after_batch`.Valid:- `draw`: time spent waiting for a batch to be drawn. Measured from `before_draw` to `before_batch`. With default prefetching settings, ideally this is as close to instantly as possible. - `batch`: total duration of all batch steps sans drawing the batch. Measured from `before_batch` to `after_batch`.- `pred`: duration of the forward pass and any additional batch modifications. Measured from `before_batch` to `after_pred`.- `loss`: duration of caculating loss. Measured from `after_pred` to `after_loss`.  ExamplesBoth examples are trained on Imagenette with an image size of 256 and batch size of 64 on a Colab P100 4CPU instance.<jupyter_code>#slow
learn = Learner(dls, xse_resnet50(n_out=dls.c), metrics=accuracy).to_fp16().profile()
learn.fit_one_cycle(2, 3e-3)<jupyter_output><empty_output><jupyter_text>When training a XResNet18, the total time spent drawing while training increases to 20 seconds, twenty percent of total fit time. This is due to Colab's four core CPU and slow disk not being able to prefetch quickly enough, causing the training process to wait on drawing a batch. In contrast in the XSEResNet50 training above, total training draw time was only 8 seconds, two percent of the total fit time.<jupyter_code>#slow
learn = Learner(dls, xresnet18(n_out=dls.c), metrics=accuracy).to_fp16().profile()
learn.fit_one_cycle(2, 3e-3)<jupyter_output><empty_output><jupyter_text>New Training LoopThe `show_training_loop` output below shows where the new `before_draw` event fits into the training loop.<jupyter_code>learn = synth_learner()
learn.show_training_loop()
#hide
#slow
from nbdev.export import notebook2script; notebook2script()<jupyter_output>Converted augment.tensor_item_tfm.ipynb.
Converted callback.cutmixup.ipynb.
Converted callback.lr_finder.ipynb.
Converted callback.simpleprofiler.ipynb.
Converted index.ipynb.
Converted schedulers.fit_flat_varied.ipynb.
{'fraction_non_alphanumeric': 0.10872848082150408, 'fraction_numerical': 0.0035034732709151315, 'mean_word_length': 3.994268476621418, 'pattern_counts': {'":': 0, '<': 40, '<?xml version=': 0, '>': 42, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '107007', 'license': 'MIT', 'n_tokens_mistral': 5095, 'n_tokens_neox': 4977, 'n_words': 1476, 'path': 'nbs/callback.simpleprofiler.ipynb', 'repo_name': 'warner-benjamin/fastai_snippets', 'size': '16555'}