{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "dca0338d", "metadata": {}, "outputs": [], "source": [ "#| hide\n", "#| eval: false\n", "! [ -e /content ] && pip install -Uqq fastai # upgrade fastai on colab" ] }, { "cell_type": "code", "execution_count": null, "id": "0e1384b5", "metadata": {}, "outputs": [], "source": [ "#| default_exp callback.training" ] }, { "cell_type": "code", "execution_count": null, "id": "1fedb5ed", "metadata": {}, "outputs": [], "source": [ "#| export\n", "from fastai.basics import *\n", "from fastai.callback.progress import *\n", "from fastai.callback.fp16 import *" ] }, { "cell_type": "code", "execution_count": null, "id": "3da31178", "metadata": {}, "outputs": [], "source": [ "#| hide\n", "from nbdev.showdoc import *\n", "from fastai.test_utils import *\n", "from fastai.vision.all import *" ] }, { "cell_type": "markdown", "id": "f826b6c7", "metadata": {}, "source": [ "# Training callbacks\n", "\n", "> Various callbacks to customize training behavior" ] }, { "cell_type": "markdown", "id": "ba6087b8", "metadata": {}, "source": [ "## ShortEpochCallback -" ] }, { "cell_type": "code", "execution_count": null, "id": "4db4681a", "metadata": {}, "outputs": [], "source": [ "#| export\n", "class ShortEpochCallback(Callback):\n", " \"Fit just `pct` of an epoch, then stop\"\n", " def __init__(self,pct=0.01,short_valid=True): self.pct,self.short_valid = pct,short_valid\n", " def after_batch(self):\n", " if self.iter/self.n_iter < self.pct: return\n", " if self.training: raise CancelTrainException\n", " if self.short_valid: raise CancelValidException" ] }, { "cell_type": "code", "execution_count": null, "id": "26a9295c", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
epochtrain_lossvalid_losstime
000:00
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "learn = synth_learner()\n", "learn.fit(1, cbs=ShortEpochCallback())" ] }, { "cell_type": "code", "execution_count": null, "id": "8d1324a4", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
epochtrain_lossvalid_losstime
08.43213500:00
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "learn = synth_learner()\n", "learn.fit(1, cbs=ShortEpochCallback(short_valid=False))" ] }, { "cell_type": "markdown", "id": "8d38bb55", "metadata": {}, "source": [ "## GradientAccumulation -" ] }, { "cell_type": "code", "execution_count": null, "id": "72e98d04", "metadata": {}, "outputs": [], "source": [ "#| export\n", "class GradientAccumulation(Callback):\n", " \"Accumulate gradients before updating weights\"\n", " order,run_valid = MixedPrecision.order-4,False\n", " def __init__(self, n_acc=32): store_attr()\n", " def before_fit(self): self.count=0\n", " def after_loss(self): self.learn.loss_grad /= self.n_acc/find_bs(self.learn.yb)\n", " def before_step(self):\n", " \"Skip weight update if we have not seen enough items\"\n", " self.learn.loss_grad *= self.n_acc/find_bs(self.learn.yb) # log correct loss\n", " self.count += find_bs(self.learn.yb)\n", " if self.count\n", " /* Turns off some styling */\n", " progress {\n", " /* gets rid of default border in Firefox and Opera. */\n", " border: none;\n", " /* Needs to be in here for Safari polyfill so background images work as expected. */\n", " background-size: auto;\n", " }\n", " .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n", " background: #F44336;\n", " }\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
epochtrain_lossvalid_losstime
00.8340620.29595000:00
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
epochtrain_lossvalid_losstime
00.8245500.29595000:00
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "#| hide\n", "class GetGrads(Callback):\n", " run_valid,order = False,GradientAccumulation.order+1\n", " def before_step(self): self.grads=to_detach(L([p.grad.clone() for p in self.model.parameters()]))\n", "\n", "def _test_acc(bs,n,cbs=None,cuda=False):\n", " with no_random(99): \n", " db=synth_dbunch(bs=bs,n_train=n,n_valid=n,cuda=cuda)\n", " learn = synth_learner(data=db,cbs=[GetGrads]+L(cbs))\n", " learn.fit(1, lr=0.01)\n", " train,valid = learn.recorder.values[-1]\n", " return train,valid,learn.get_grads.grads\n", "\n", "acc_cb = GradientAccumulation(n_acc=8)\n", "\n", "train1,valid1,grads1 = _test_acc(8,1)\n", "train2,valid2,grads2 = _test_acc(1,8,acc_cb)\n", "\n", "#grads should be same, valid loss same, train loss different\n", "test_close(grads2,grads1)\n", "test_close(valid2, valid1)\n", "test_ne(train2, train1)" ] }, { "cell_type": "code", "execution_count": null, "id": "008106b9", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
epochtrain_lossvalid_losstime
00.8340620.29595000:00
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
epochtrain_lossvalid_losstime
00.8245500.29595000:00
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "#| hide\n", "#| cuda\n", "fp16_cb = MixedPrecision(init_scale=1024)\n", "train1,valid1,grads1 = _test_acc(8,1, fp16_cb, cuda=True)\n", "train2,valid2,grads2 = _test_acc(1,8, [acc_cb,fp16_cb], cuda=True)\n", "test_close(grads2,grads1, eps=0.01)\n", "test_close(valid2, valid1)\n", "test_ne(train2, train1)" ] }, { "cell_type": "markdown", "id": "405adaab", "metadata": {}, "source": [ "When the number of steps per accumulation is higher than the number of batches, the parameters (and therefore validation loss) don't change at all:" ] }, { "cell_type": "code", "execution_count": null, "id": "6f302106", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
epochtrain_lossvalid_losstime
020.98755826.84948000:00
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "learn = synth_learner()\n", "learn.fit(1, lr=0.01, cbs=GradientAccumulation(n_acc=1000))\n", "# ensure valid_loss didn't change\n", "assert learn.recorder.values[-1][1] == learn.recorder.values[0][1]" ] }, { "cell_type": "markdown", "id": "4273e63b", "metadata": {}, "source": [ "## GradientClip -" ] }, { "cell_type": "code", "execution_count": null, "id": "3b53eb46", "metadata": {}, "outputs": [], "source": [ "#| export\n", "class GradientClip(Callback):\n", " \"Clip norm of gradients\"\n", " order=MixedPrecision.order+1\n", " def __init__(self,max_norm:float=1., norm_type:float=2.0): store_attr()\n", " def before_step(self): nn.utils.clip_grad_norm_(self.parameters(), self.max_norm, self.norm_type)" ] }, { "cell_type": "markdown", "id": "38cb1bfc", "metadata": {}, "source": [ "Normally if we use a learning rate that is too high, our training will diverge. This even happens if we use mixed precision training, which avoid infinities by using dynamic loss scaling, but still diverges:" ] }, { "cell_type": "code", "execution_count": null, "id": "f187a26a", "metadata": {}, "outputs": [], "source": [ "fp16 = MixedPrecision()" ] }, { "cell_type": "code", "execution_count": null, "id": "d6845d34", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
epochtrain_lossvalid_losstime
038.21413825.26900500:00
1377.145508890.01037600:00
2839.3928839965.74707000:00
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "set_seed(99)\n", "learn = synth_learner(lr=1.1, cuda=True)\n", "learn.fit(3, cbs=fp16)" ] }, { "cell_type": "markdown", "id": "93fc8b5e", "metadata": {}, "source": [ "By adding the `GradientClip` callback, the gradient `norm_type` (default:2) norm is clipped to at most `max_norm` (default:1) using `nn.utils.clip_grad_norm_`, which can avoid loss divergence:" ] }, { "cell_type": "code", "execution_count": null, "id": "00622e32", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
epochtrain_lossvalid_losstime
02.0394282.37217700:00
11.4024250.30072800:00
21.0135480.33261000:00
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "set_seed(99)\n", "learn = synth_learner(lr=1.1, cuda=True)\n", "learn.fit(3, cbs=[GradientClip,fp16])" ] }, { "cell_type": "markdown", "id": "26b5c825", "metadata": {}, "source": [ "## BnFreeze" ] }, { "cell_type": "code", "execution_count": null, "id": "11915cc7", "metadata": {}, "outputs": [], "source": [ "#| export\n", "bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)\n", "\n", "def set_bn_eval(m:nn.Module, use_eval=True)->None:\n", " \"Set bn layers in eval mode for all recursive children of `m`.\"\n", " for l in m.children():\n", " if isinstance(l, bn_types) and not next(l.parameters()).requires_grad:\n", " if use_eval: l.eval()\n", " else: l.train()\n", " set_bn_eval(l)\n", "\n", "class BnFreeze(Callback):\n", " run_after=TrainEvalCallback\n", " \"Freeze moving average statistics in all non-trainable batchnorm layers.\"\n", " def before_train(self):\n", " set_bn_eval(self.model)" ] }, { "cell_type": "markdown", "id": "f4cc6f33", "metadata": {}, "source": [ "`BnFreeze` is useful when you'd like to train two separate models that have a common feature extractor / body. The only part of the model that's different is the head that you attach for transfer learning.
\n", "\n", "`Learner.freeze()` doesn't suffice here as the `BatchNorm` layers are trainable by default, and running mean and std of batches are tracked. For feature extractors to fully match, you need to set `train_bn=False` and these stats need to be frozen as well, which is precisely the function of `BnFreeze`." ] }, { "cell_type": "code", "execution_count": null, "id": "7e7ecaf1", "metadata": {}, "outputs": [], "source": [ "#| slow\n", "path = untar_data(URLs.MNIST_TINY)\n", "dls = ImageDataLoaders.from_folder(path, valid_pct=0.2)" ] }, { "cell_type": "markdown", "id": "a0dd6ca3", "metadata": {}, "source": [ "https://pytorch.org/tutorials/intermediate/memory_format_tutorial.htmlWe first demonstrate the mismatch of the running stats when using only `train_bn=False`, by creating a `Learner`...:" ] }, { "cell_type": "code", "execution_count": null, "id": "5797dce8", "metadata": {}, "outputs": [], "source": [ "#| slow\n", "learn1 = vision_learner(deepcopy(dls), resnet18, pretrained=True, train_bn=False)" ] }, { "cell_type": "markdown", "id": "5f5c4962", "metadata": {}, "source": [ "...and grab the first `BatchNorm` layer, and store its running mean: " ] }, { "cell_type": "code", "execution_count": null, "id": "5c74c25a", "metadata": {}, "outputs": [], "source": [ "#| slow\n", "m = learn1.model[0][1].running_mean.clone()" ] }, { "cell_type": "markdown", "id": "5421ba7c", "metadata": {}, "source": [ "You can see that now that running mean has changed:" ] }, { "cell_type": "code", "execution_count": null, "id": "c2b69807", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
epochtrain_lossvalid_losstime
01.1483030.73940400:12
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "#| slow\n", "learn1.fit(1, lr=0.02)\n", "test_ne(to_detach(learn1.model[0][1].running_mean), m)" ] }, { "cell_type": "markdown", "id": "42fad84e", "metadata": {}, "source": [ "When we use the `BnFreeze` callback, the running statistics will not be changed during training. This is often important for getting good results from transfer learning." ] }, { "cell_type": "code", "execution_count": null, "id": "d2999c5d", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
epochtrain_lossvalid_losstime
00.4785940.27077200:10
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "#| slow\n", "learn1 = vision_learner(deepcopy(dls), resnet18, pretrained=True, train_bn=False, cbs=BnFreeze)\n", "m = learn1.model[0][1].running_mean.detach().clone()\n", "learn1.fit(1, lr=0.02)\n", "test_eq(to_detach(learn1.model[0][1].running_mean), m)" ] }, { "cell_type": "markdown", "id": "b8d5320a", "metadata": {}, "source": [ "## Export -" ] }, { "cell_type": "code", "execution_count": null, "id": "b8f905bb", "metadata": {}, "outputs": [], "source": [ "#| hide\n", "from nbdev import nbdev_export\n", "nbdev_export()" ] } ], "metadata": { "jupytext": { "split_at_heading": true }, "kernelspec": { "display_name": "python3", "language": "python", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 5 }