fastai--fastai
wehub-resource-sync
1460 行
51 KiB
Plaintext
1460 行
51 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "833a691f",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| eval: false\n",
|
|
"! [ -e /content ] && pip install -Uqq fastai # upgrade fastai on colab"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "raw",
|
|
"id": "be0957fc",
|
|
"metadata": {},
|
|
"source": [
|
|
"---\n",
|
|
"skip_exec: true\n",
|
|
"---"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "8d2f192e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"from fastai.basics import *\n",
|
|
"from fastai.callback.progress import *\n",
|
|
"\n",
|
|
"from torch.amp import GradScaler,autocast\n",
|
|
"from torch.amp.grad_scaler import OptState"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "46e90a8e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| default_exp callback.fp16"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f25b19cd",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"from fastai.test_utils import *\n",
|
|
"from nbdev.showdoc import *"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "e6078403",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Mixed precision training\n",
|
|
"\n",
|
|
"> Callback and utility functions to allow mixed precision training "
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "0d7f069a",
|
|
"metadata": {},
|
|
"source": [
|
|
"## A little bit of theory"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "7b49c955",
|
|
"metadata": {},
|
|
"source": [
|
|
"A very nice and clear introduction to mixed precision training is [this video from NVIDIA](https://on-demand.gputechconf.com/gtc/2019/video/_/S9143/)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "82e2378e",
|
|
"metadata": {},
|
|
"source": [
|
|
"### What's half precision?"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "59bc3f53",
|
|
"metadata": {},
|
|
"source": [
|
|
"In neural nets, all the computations are usually done in single precision, which means all the floats in all the arrays that represent inputs, activations, weights... are 32-bit floats (FP32 in the rest of this post). An idea to reduce memory usage (and avoid those annoying cuda errors) has been to try and do the same thing in half-precision, which means using 16-bits floats (or FP16 in the rest of this post). By definition, they take half the space in RAM, and in theory could allow you to double the size of your model and double your batch size.\n",
|
|
"\n",
|
|
"Another very nice feature is that NVIDIA developed its latest GPUs (the Volta generation) to take fully advantage of half-precision tensors. Basically, if you give half-precision tensors to those, they'll stack them so that each core can do more operations at the same time, and theoretically gives an 8x speed-up (sadly, just in theory).\n",
|
|
"\n",
|
|
"So training at half precision is better for your memory usage, way faster if you have a Volta GPU (still a tiny bit faster if you don't since the computations are easiest). How do we do it? Super easily in pytorch, we just have to put .half() everywhere: on the inputs of our model and all the parameters. Problem is that you usually won't see the same accuracy in the end (so it happens sometimes) because half-precision is... well... not as precise ;)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "14a028e8",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Problems with half-precision:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "f9a84441",
|
|
"metadata": {},
|
|
"source": [
|
|
"To understand the problems with half precision, let's look briefly at what an FP16 looks like (more information [here](https://en.wikipedia.org/wiki/Half-precision_floating-point_format)).\n",
|
|
"\n",
|
|
"\n",
|
|
"\n",
|
|
"The sign bit gives us +1 or -1, then we have 5 bits to code an exponent between -14 and 15, while the fraction part has the remaining 10 bits. Compared to FP32, we have a smaller range of possible values (2e-14 to 2e15 roughly, compared to 2e-126 to 2e127 for FP32) but also a smaller *offset*.\n",
|
|
"\n",
|
|
"For instance, between 1 and 2, the FP16 format only represents the number 1, 1+2e-10, 1+2*2e-10... which means that 1 + 0.0001 = 1 in half precision. That's what will cause a certain numbers of problems, specifically three that can occur and mess up your training.\n",
|
|
"\n",
|
|
"1. The weight update is imprecise: inside your optimizer, you basically do w = w - lr * w.grad for each weight of your network. The problem in performing this operation in half precision is that very often, w.grad is several orders of magnitude below w, and the learning rate is also small. The situation where w=1 and lr*w.grad is 0.0001 (or lower) is therefore very common, but the update doesn't do anything in those cases.\n",
|
|
"\n",
|
|
"2. Your gradients can underflow. In FP16, your gradients can easily be replaced by 0 because they are too low.\n",
|
|
"\n",
|
|
"3. Your activations or loss can overflow. The opposite problem from the gradients: it's easier to hit nan (or infinity) in FP16 precision, and your training might more easily diverge."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "595e3851",
|
|
"metadata": {},
|
|
"source": [
|
|
"### The solution: mixed precision training"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "63380f3f",
|
|
"metadata": {},
|
|
"source": [
|
|
"To address those three problems, we don't fully train in FP16 precision. As the name mixed training implies, some of the operations will be done in FP16, others in FP32. This is mainly to take care of the first problem listed above. For the next two there are additional tricks.\n",
|
|
"\n",
|
|
"The main idea is that we want to do the forward pass and the gradient computation in half precision (to go fast) but the update in single precision (to be more precise). It's okay if w and grad are both half floats, but when we do the operation w = w - lr * grad, we need to compute it in FP32. That way our 1 + 0.0001 is going to be 1.0001. \n",
|
|
"\n",
|
|
"This is why we keep a copy of the weights in FP32 (called master model). Then, our training loop will look like:\n",
|
|
"\n",
|
|
"1. compute the output with the FP16 model, then the loss\n",
|
|
"2. back-propagate the gradients in half-precision.\n",
|
|
"3. copy the gradients in FP32 precision\n",
|
|
"4. do the update on the master model (in FP32 precision)\n",
|
|
"5. copy the master model in the FP16 model.\n",
|
|
"\n",
|
|
"Note that we lose precision during step 5, and that the 1.0001 in one of the weights will go back to 1. But if the next update corresponds to add 0.0001 again, since the optimizer step is done on the master model, the 1.0001 will become 1.0002 and if we eventually go like this up to 1.0005, the FP16 model will be able to tell the difference.\n",
|
|
"\n",
|
|
"That takes care of problem 1. For the second problem, we use something called gradient scaling: to avoid the gradients getting zeroed by the FP16 precision, we multiply the loss by a scale factor (scale=512 for instance). That way we can push the gradients to the right in the next figure, and have them not become zero.\n",
|
|
"\n",
|
|
"\n",
|
|
"\n",
|
|
"Of course we don't want those 512-scaled gradients to be in the weight update, so after converting them into FP32, we can divide them by this scale factor (once they have no risks of becoming 0). This changes the loop to:\n",
|
|
"\n",
|
|
"1. compute the output with the FP16 model, then the loss.\n",
|
|
"2. multiply the loss by scale then back-propagate the gradients in half-precision.\n",
|
|
"3. copy the gradients in FP32 precision then divide them by scale.\n",
|
|
"4. do the update on the master model (in FP32 precision).\n",
|
|
"5. copy the master model in the FP16 model.\n",
|
|
"\n",
|
|
"For the last problem, the tricks offered by NVIDIA are to leave the batchnorm layers in single precision (they don't have many weights so it's not a big memory challenge) and compute the loss in single precision (which means converting the last output of the model in single precision before passing it to the loss).\n",
|
|
"\n",
|
|
""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "a639f0af",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Dynamic loss scaling"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "7bcf457e",
|
|
"metadata": {},
|
|
"source": [
|
|
"The only annoying thing with the previous implementation of mixed precision training is that it introduces one new hyper-parameter to tune, the value of the loss scaling. Fortunately for us, there is a way around this. We want the loss scaling to be as high as possible so that our gradients can use the whole range of representation, so let's first try a really high value. In all likelihood, this will cause our gradients or our loss to overflow, and we will try again with half that big value, and again, until we get to the largest loss scale possible that doesn't make our gradients overflow.\n",
|
|
"\n",
|
|
"This value will be perfectly fitted to our model and can continue to be dynamically adjusted as the training goes, if it's still too high, by just halving it each time we overflow. After a while though, training will converge and gradients will start to get smaller, so we al\n",
|
|
"so need a mechanism to get this dynamic loss scale larger if it's safe to do so. The strategy used in the Apex library is to multiply the loss scale by 2 each time we had a given number of iterations without overflowing."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "e6a36f48",
|
|
"metadata": {},
|
|
"source": [
|
|
"### BFloat16 Mixed Precision"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "2df51c49",
|
|
"metadata": {},
|
|
"source": [
|
|
"BFloat16 (BF16) is 16-bit floating point format developed by Google Brain. BF16 has the same exponent as FP32 leaving 7-bits for the fraction. This gives BF16 the same range as FP32, but significantly less precision.\n",
|
|
"\n",
|
|
"Since it has same range as FP32, BF16 Mixed Precision training skips the scaling steps. All other Mixed Precision steps remain the same as FP16 Mixed Precision.\n",
|
|
"\n",
|
|
"BF16 Mixed Precision requires Ampere or newer hardware. Not all PyTorch operations are supported.\n",
|
|
"\n",
|
|
"To train in BF16 Mixed Precision pass `amp_mode=AMPMode.BF16` or `amp_mode='bf16'` to `MixedPrecision`, or use the `Learner.to_bf16` convenience method."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "22b33bc9",
|
|
"metadata": {},
|
|
"source": [
|
|
"## MixedPrecision -"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "2739b62c",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"class AMPMode(Enum):\n",
|
|
" \"Automatic mixed precision modes for ease of completion\"\n",
|
|
" FP16 = 'fp16'\n",
|
|
" BF16 = 'bf16'"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3d30daa9",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@delegates(GradScaler)\n",
|
|
"class MixedPrecision(Callback):\n",
|
|
" \"Mixed precision training using Pytorch's Automatic Mixed Precision (AMP)\"\n",
|
|
" order = 10\n",
|
|
" def __init__(self,\n",
|
|
" amp_mode:str|AMPMode=AMPMode.FP16, # Mixed Precision training mode. Supports fp16 and bf16.\n",
|
|
" **kwargs\n",
|
|
" ):\n",
|
|
" amp_mode = AMPMode(amp_mode)\n",
|
|
" store_attr(names='amp_mode')\n",
|
|
" self.kwargs = kwargs\n",
|
|
"\n",
|
|
" def before_fit(self):\n",
|
|
" if self.amp_mode == AMPMode.BF16:\n",
|
|
" if torch.cuda.is_available() and not torch.cuda.is_bf16_supported():\n",
|
|
" raise ValueError(\"Unsupported GPU for bfloat16 mixed precision training\")\n",
|
|
" dtype = torch.bfloat16\n",
|
|
" elif self.amp_mode == AMPMode.FP16:\n",
|
|
" dtype = torch.float16\n",
|
|
" else:\n",
|
|
" raise ValueError(f\"Unrecognized precision: {self.amp_mode}\")\n",
|
|
" # `GradScaler` is not needed for bfloat16 as fp32 and bf16 have the same range\n",
|
|
" self.kwargs['enabled'] = dtype == torch.float16\n",
|
|
" self.autocast,self.learn.scaler,self.scales = autocast('cuda', dtype=dtype),GradScaler('cuda', **self.kwargs),L()\n",
|
|
"\n",
|
|
" def before_batch(self): self.autocast.__enter__()\n",
|
|
" def after_pred(self):\n",
|
|
" self.learn.pred = to_float(self.pred)\n",
|
|
" def after_loss(self): self.autocast.__exit__(None, None, None)\n",
|
|
" def before_backward(self): self.learn.loss_grad = self.scaler.scale(self.loss_grad)\n",
|
|
" def before_step(self):\n",
|
|
" \"Use `self` as a fake optimizer. `self.skipped` will be set to True `after_step` if gradients overflow.\"\n",
|
|
" self.skipped=True\n",
|
|
" self.scaler.step(self)\n",
|
|
" if self.skipped: raise CancelStepException()\n",
|
|
" self.scales.append(self.scaler.get_scale())\n",
|
|
" def after_step(self): self.learn.scaler.update()\n",
|
|
" def after_fit(self): self.autocast,self.learn.scaler,self.scales = None,None,None\n",
|
|
"\n",
|
|
" @property\n",
|
|
" def param_groups(self):\n",
|
|
" \"Pretend to be an optimizer for `GradScaler`\"\n",
|
|
" return self.opt.param_groups\n",
|
|
" def step(self, *args, **kwargs):\n",
|
|
" \"Fake optimizer step to detect whether this batch was skipped from `GradScaler`\"\n",
|
|
" self.skipped=False"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "67797261",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/markdown": [
|
|
"---\n",
|
|
"\n",
|
|
"[source](https://github.com/fastai/fastai/blob/main/fastai/callback/fp16.py#L25){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
|
|
"\n",
|
|
"### MixedPrecision\n",
|
|
"\n",
|
|
"> MixedPrecision (amp_mode:str|AMPMode=<AMPMode.FP16: 'fp16'>,\n",
|
|
"> init_scale:float=65536.0, growth_factor:float=2.0,\n",
|
|
"> backoff_factor:float=0.5, growth_interval:int=2000,\n",
|
|
"> enabled:bool=True)\n",
|
|
"\n",
|
|
"*Mixed precision training using Pytorch's Automatic Mixed Precision (AMP)*\n",
|
|
"\n",
|
|
"| | **Type** | **Default** | **Details** |\n",
|
|
"| -- | -------- | ----------- | ----------- |\n",
|
|
"| amp_mode | str \\| __main__.AMPMode | AMPMode.FP16 | Mixed Precision training mode. Supports fp16 and bf16. |\n",
|
|
"| init_scale | float | 65536.0 | |\n",
|
|
"| growth_factor | float | 2.0 | |\n",
|
|
"| backoff_factor | float | 0.5 | |\n",
|
|
"| growth_interval | int | 2000 | |\n",
|
|
"| enabled | bool | True | |"
|
|
],
|
|
"text/plain": [
|
|
"---\n",
|
|
"\n",
|
|
"[source](https://github.com/fastai/fastai/blob/main/fastai/callback/fp16.py#L25){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
|
|
"\n",
|
|
"### MixedPrecision\n",
|
|
"\n",
|
|
"> MixedPrecision (amp_mode:str|AMPMode=<AMPMode.FP16: 'fp16'>,\n",
|
|
"> init_scale:float=65536.0, growth_factor:float=2.0,\n",
|
|
"> backoff_factor:float=0.5, growth_interval:int=2000,\n",
|
|
"> enabled:bool=True)\n",
|
|
"\n",
|
|
"*Mixed precision training using Pytorch's Automatic Mixed Precision (AMP)*\n",
|
|
"\n",
|
|
"| | **Type** | **Default** | **Details** |\n",
|
|
"| -- | -------- | ----------- | ----------- |\n",
|
|
"| amp_mode | str \\| __main__.AMPMode | AMPMode.FP16 | Mixed Precision training mode. Supports fp16 and bf16. |\n",
|
|
"| init_scale | float | 65536.0 | |\n",
|
|
"| growth_factor | float | 2.0 | |\n",
|
|
"| backoff_factor | float | 0.5 | |\n",
|
|
"| growth_interval | int | 2000 | |\n",
|
|
"| enabled | bool | True | |"
|
|
]
|
|
},
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"show_doc(MixedPrecision)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "6dccb096",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"class FP16TestCallback(Callback):\n",
|
|
" \"Asserts that predictions are `float16` values\"\n",
|
|
" order = 9\n",
|
|
" def after_pred(self):\n",
|
|
" assert listify(flatten(self.pred))[0].dtype==torch.float16"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "c834ad9e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"class BF16TestCallback(Callback):\n",
|
|
" \"Asserts that predictions are `bfloat16` values\"\n",
|
|
" order = 9\n",
|
|
" def after_pred(self):\n",
|
|
" assert listify(flatten(self.pred))[0].dtype==torch.bfloat16"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "c252a5cd",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/html": [
|
|
"<table border=\"1\" class=\"dataframe\">\n",
|
|
" <thead>\n",
|
|
" <tr style=\"text-align: left;\">\n",
|
|
" <th>epoch</th>\n",
|
|
" <th>train_loss</th>\n",
|
|
" <th>valid_loss</th>\n",
|
|
" <th>time</th>\n",
|
|
" </tr>\n",
|
|
" </thead>\n",
|
|
" <tbody>\n",
|
|
" <tr>\n",
|
|
" <td>0</td>\n",
|
|
" <td>17.554865</td>\n",
|
|
" <td>14.357819</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>1</td>\n",
|
|
" <td>17.006779</td>\n",
|
|
" <td>13.436550</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>2</td>\n",
|
|
" <td>16.414442</td>\n",
|
|
" <td>12.542552</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" </tbody>\n",
|
|
"</table>"
|
|
],
|
|
"text/plain": [
|
|
"<IPython.core.display.HTML object>"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
}
|
|
],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| cuda\n",
|
|
"set_seed(99, True)\n",
|
|
"learn = synth_learner(cbs=[MixedPrecision,FP16TestCallback], cuda=True)\n",
|
|
"learn.model = nn.Sequential(nn.Linear(1,1), nn.Linear(1,1)).cuda()\n",
|
|
"learn.opt_func = partial(SGD, mom=0.)\n",
|
|
"learn.splitter = lambda m: [list(m[0].parameters()), list(m[1].parameters())]\n",
|
|
"learn.fit(3)\n",
|
|
"assert learn.recorder.values[-1][-1]<learn.recorder.values[0][-1]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d28b3f6a",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/html": [
|
|
"<table border=\"1\" class=\"dataframe\">\n",
|
|
" <thead>\n",
|
|
" <tr style=\"text-align: left;\">\n",
|
|
" <th>epoch</th>\n",
|
|
" <th>train_loss</th>\n",
|
|
" <th>valid_loss</th>\n",
|
|
" <th>time</th>\n",
|
|
" </tr>\n",
|
|
" </thead>\n",
|
|
" <tbody>\n",
|
|
" <tr>\n",
|
|
" <td>0</td>\n",
|
|
" <td>87.652245</td>\n",
|
|
" <td>72.425194</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>1</td>\n",
|
|
" <td>86.457306</td>\n",
|
|
" <td>70.571136</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>2</td>\n",
|
|
" <td>85.303947</td>\n",
|
|
" <td>68.533089</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" </tbody>\n",
|
|
"</table>"
|
|
],
|
|
"text/plain": [
|
|
"<IPython.core.display.HTML object>"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
}
|
|
],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| cuda\n",
|
|
"#Multioutput version\n",
|
|
"set_seed(99, True)\n",
|
|
"learn = synth_learner(cbs=[MixedPrecision,FP16TestCallback], cuda=True)\n",
|
|
"class MultiOutputModel(Module):\n",
|
|
" def __init__(self): self.linear1, self.linear2 = nn.Linear(1,1) , nn.Linear(1,1)\n",
|
|
" def forward(self,x): return self.linear1(x), self.linear2(x)\n",
|
|
"def multioutputloss(pred, val): return ((val-pred[0]).abs() + 0.5 * (val-pred[1]).abs()).sum()\n",
|
|
"learn.model = MultiOutputModel()\n",
|
|
"learn.opt_func = partial(SGD, mom=0.)\n",
|
|
"learn.splitter = lambda m: [list(m.linear1.parameters()), list(m.linear2.parameters())]\n",
|
|
"learn.loss_func=multioutputloss\n",
|
|
"learn.fit(3)\n",
|
|
"assert learn.recorder.values[-1][-1]<learn.recorder.values[0][-1]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "e713f25e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| cuda\n",
|
|
"if torch.cuda.is_bf16_supported():\n",
|
|
" set_seed(99, True)\n",
|
|
" learn = synth_learner(cbs=[MixedPrecision(amp_mode=AMPMode.BF16),BF16TestCallback], cuda=True)\n",
|
|
" learn.model = nn.Sequential(nn.Linear(1,1), nn.Linear(1,1)).cuda()\n",
|
|
" learn.opt_func = partial(SGD, mom=0.)\n",
|
|
" learn.splitter = lambda m: [list(m[0].parameters()), list(m[1].parameters())]\n",
|
|
" learn.fit(3)\n",
|
|
" assert learn.recorder.values[-1][-1]<learn.recorder.values[0][-1]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d56db335",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"@delegates(GradScaler)\n",
|
|
"def to_fp16(self:Learner, **kwargs):\n",
|
|
" \"Set `Learner` to float16 mixed precision using PyTorch AMP\"\n",
|
|
" return self.add_cb(MixedPrecision(**kwargs))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "6c6c0e6d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"def to_bf16(self:Learner):\n",
|
|
" \"Set `Learner` to bfloat16 mixed precision using PyTorch AMP\"\n",
|
|
" return self.add_cb(MixedPrecision(amp_mode=AMPMode.BF16))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3d78320f",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"def to_fp32(self:Learner):\n",
|
|
" \"Set `Learner` to float32 precision\"\n",
|
|
" return self.remove_cb(MixedPrecision)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "05a6d649",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Util functions"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "c7b6c0af",
|
|
"metadata": {},
|
|
"source": [
|
|
"Before going in the main `Callback` we will need some helper functions. We use the ones from the [APEX library](https://github.com/NVIDIA/apex)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3a800e16",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export \n",
|
|
"from fastai.fp16_utils import convert_network, model_grads_to_master_grads, master_params_to_model_params"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "4dc907e9",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Converting the model to FP16"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "a39531d4",
|
|
"metadata": {},
|
|
"source": [
|
|
"We will need a function to convert all the layers of the model to FP16 precision except the BatchNorm-like layers (since those need to be done in FP32 precision to be stable). In Apex, the function that does this for us is `convert_network`. We can use it to put the model in FP16 or back to FP32."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1cd78ab5",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"model = nn.Sequential(nn.Linear(10,30), nn.BatchNorm1d(30), nn.Linear(30,2)).cuda()\n",
|
|
"model = convert_network(model, torch.float16)\n",
|
|
"\n",
|
|
"for i,t in enumerate([torch.float16, torch.float32, torch.float16]):\n",
|
|
" test_eq(model[i].weight.dtype, t)\n",
|
|
" test_eq(model[i].bias.dtype, t)\n",
|
|
" \n",
|
|
"model = nn.Sequential(nn.Linear(10,30), BatchNorm(30, ndim=1), nn.Linear(30,2)).cuda()\n",
|
|
"model = convert_network(model, torch.float16)\n",
|
|
"\n",
|
|
"for i,t in enumerate([torch.float16, torch.float32, torch.float16]):\n",
|
|
" test_eq(model[i].weight.dtype, t)\n",
|
|
" test_eq(model[i].bias.dtype, t)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "7cdcf4c4",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Creating the master copy of the parameters"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "2919a018",
|
|
"metadata": {},
|
|
"source": [
|
|
"From our model parameters (mostly in FP16), we'll want to create a copy in FP32 (master parameters) that we will use for the step in the optimizer. Optionally, we concatenate all the parameters to do one flat big tensor, which can make that step a little bit faster.\n",
|
|
"\n",
|
|
"We can't use the FP16 util function here as it doesn't handle multiple parameter groups, which is the thing we use to:\n",
|
|
"\n",
|
|
"- do transfer learning and freeze some layers\n",
|
|
"- apply discriminative learning rates\n",
|
|
"- don't apply weight decay to some layers (like BatchNorm) or the bias terms"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "73fd0b34",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"from torch.nn.utils import parameters_to_vector"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1e983fc9",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"def get_master(\n",
|
|
" opt:Optimizer, # Optimizer from which to retrieve model params\n",
|
|
" flat_master:bool=False, # Flatten fp32 params into a vector for better performance\n",
|
|
") -> list: # List of fp16 params, and list of fp32 params\n",
|
|
" \"Creates fp16 model params given an initialized `Optimizer`, also returning fp32 model params. \"\n",
|
|
" model_params = [[param for param in pg if getattr(param, 'requires_grad', False) and hasattr(param, 'data')] for pg in opt.param_lists]\n",
|
|
" if flat_master:\n",
|
|
" master_params = []\n",
|
|
" for pg in model_params:\n",
|
|
" mp = parameters_to_vector([param.data.float() for param in pg])\n",
|
|
" mp = nn.Parameter(mp, requires_grad=True)\n",
|
|
" if mp.grad is None: mp.grad = mp.new(*mp.size())\n",
|
|
" master_params.append([mp])\n",
|
|
" else:\n",
|
|
" master_params = [[nn.Parameter(param.data.clone().float().detach(), requires_grad=True) for param in pg] for pg in model_params]\n",
|
|
" return model_params, master_params"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "4cc4977e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| cuda\n",
|
|
"learn = synth_learner()\n",
|
|
"learn.model = convert_network(nn.Sequential(nn.Linear(1,1), nn.Linear(1,1)), torch.float16).cuda()\n",
|
|
"learn.splitter = lambda m: [list(m[0].parameters()), list(m[1].parameters())]\n",
|
|
"learn.opt = learn.opt_func(learn.splitter(learn.model), learn.lr)\n",
|
|
"model_p,master_p = get_master(learn.opt)\n",
|
|
"test_eq(len(model_p), 2) #2 pqrqm groups\n",
|
|
"test_eq(len(master_p), 2)\n",
|
|
"for pg1,pg2 in zip(model_p,master_p):\n",
|
|
" test_eq([p.float() for p in pg1], pg2) #Same values but different types\n",
|
|
" for p in pg1: assert p.dtype == torch.float16"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f308ba72",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| cuda\n",
|
|
"#Flattened version\n",
|
|
"model_pf,master_pf = get_master(learn.opt, flat_master=True)\n",
|
|
"test_eq(len(model_pf), 2) #2 pqrqm groups\n",
|
|
"test_eq(len(master_pf), 2)\n",
|
|
"for pg1,pg2 in zip(model_pf,master_pf):\n",
|
|
" test_eq(len(pg2), 1) #One flattened tensor\n",
|
|
" test_eq([p.float().squeeze() for p in pg1], [p for p in pg2[0]]) #Same values but different types\n",
|
|
" for p in pg1: assert p.dtype == torch.float16"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "92d06e9d",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Copy the gradients from model params to master params"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "448551f7",
|
|
"metadata": {},
|
|
"source": [
|
|
"After the backward pass, all gradients must be copied to the master params before the optimizer step can be done in FP32. The corresponding function in the Apex utils is `model_grads_to_master_grads` but we need to adapt it to work with param groups."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "84d2c17a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export \n",
|
|
"def to_master_grads( \n",
|
|
" model_pgs:list, # Fp16 model parameters to copy gradients from\n",
|
|
" master_pgs:list, # Fp32 model parameters to copy gradients to\n",
|
|
" flat_master:bool=False, # Whether or not fp32 parameters were previously flattened\n",
|
|
"):\n",
|
|
" \"Move fp16 model gradients to fp32 master gradients\"\n",
|
|
" for (model_params,master_params) in zip(model_pgs,master_pgs):\n",
|
|
" model_grads_to_master_grads(model_params, master_params, flat_master=flat_master)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "0b2abd67",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"torch.Size([16, 1])"
|
|
]
|
|
},
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| cuda\n",
|
|
"xb,yb = learn.dls.one_batch()\n",
|
|
"pred = learn.model.cuda()(xb.cuda().half())\n",
|
|
"loss = F.mse_loss(pred, yb.cuda().half())\n",
|
|
"loss.backward()\n",
|
|
"to_master_grads(model_p, master_p)\n",
|
|
"to_master_grads(model_pf, master_pf, flat_master=True)\n",
|
|
"test_eq([[p.grad.float() for p in pg] for pg in model_p],\n",
|
|
" [[p.grad for p in pg] for pg in master_p])\n",
|
|
"test_eq([[p.grad.float().squeeze() for p in pg] for pg in model_pf], \n",
|
|
" [[p for p in pg[0].grad] for pg in master_pf])\n",
|
|
"xb.shape"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "583e7f20",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Copy the master params to the model params"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "cc22fb5c",
|
|
"metadata": {},
|
|
"source": [
|
|
"After the step, we need to copy back the master parameters to the model parameters for the next update. The corresponding function in Apex is `master_params_to_model_params`."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d11afa3d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export \n",
|
|
"def to_model_params(\n",
|
|
" model_pgs:list, # Fp16 model params to copy to\n",
|
|
" master_pgs:list, # Fp32 master params to copy from\n",
|
|
" flat_master:bool=False # Whether master_pgs was previously flattened\n",
|
|
")->None:\n",
|
|
" \"Copy updated fp32 master params to fp16 model params after gradient step. \" \n",
|
|
" for (model_params,master_params) in zip(model_pgs,master_pgs):\n",
|
|
" master_params_to_model_params(model_params, master_params, flat_master=flat_master)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "a68ed4e2",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| cuda\n",
|
|
"learn.opt.params = master_p\n",
|
|
"learn.opt.step()\n",
|
|
"to_model_params(model_p, master_p)\n",
|
|
"test_close([p.float() for pg in model_p for p in pg], [p for pg in master_p for p in pg], eps=1e-3)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d7cadca5",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| cuda\n",
|
|
"learn.opt.params = master_pf\n",
|
|
"learn.opt.step()\n",
|
|
"to_model_params(model_pf, master_pf, flat_master=True)\n",
|
|
"test_close([p.float().squeeze() for pg in model_pf for p in pg], [p for pg in master_pf for p in pg[0]], eps=1e-3)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "ce43454a",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Checking for overflow"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "0bcc3a94",
|
|
"metadata": {},
|
|
"source": [
|
|
"For dynamic loss scaling, we need to know when the gradients have gone up to infinity. It's faster to check it on the sum than to do `torch.isinf(x).any()`."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "4301268d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export \n",
|
|
"def test_overflow(x:torch.Tensor):\n",
|
|
" \"Tests whether fp16 gradients have overflown.\"\n",
|
|
" s = float(x.float().sum())\n",
|
|
" return (s == float('inf') or s == float('-inf') or s != s)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "648b9a00",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"x = torch.randn(3,4)\n",
|
|
"assert not test_overflow(x)\n",
|
|
"x[1,2] = float('inf')\n",
|
|
"assert test_overflow(x)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "bf6cd167",
|
|
"metadata": {},
|
|
"source": [
|
|
"Then we can use it in the following function that checks for gradient overflow:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f2e64c61",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export \n",
|
|
"def grad_overflow(pgs:list)->bool: \n",
|
|
" \"Tests all fp16 parameters in pgs for gradient overflow\"\n",
|
|
" for pg in pgs:\n",
|
|
" for p in pg:\n",
|
|
" if p.grad is not None and test_overflow(p.grad.data): return True\n",
|
|
" return False"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "e4942e58",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| cuda\n",
|
|
"assert not grad_overflow(model_p)\n",
|
|
"assert not grad_overflow(model_pf)\n",
|
|
"model_p[1][0].grad.data[0,0] = float('inf')\n",
|
|
"model_pf[0][1].grad.data[0] = float('inf')\n",
|
|
"assert grad_overflow(model_p)\n",
|
|
"assert grad_overflow(model_pf)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "442ea522",
|
|
"metadata": {},
|
|
"source": [
|
|
"## NonNativeMixedPrecision -"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "65bc0f15",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"def copy_clone(d):\n",
|
|
" return {k:(v.detach().clone().float() if isinstance(v,Tensor) else v) for k,v in d.items()}"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "a35f40af",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"def _copy_state(opt, pgs1, pgs2):\n",
|
|
" opt.param_lists = pgs2\n",
|
|
" for pg1,pg2 in zip(pgs1, pgs2):\n",
|
|
" for p1,p2 in zip(pg1, pg2): opt.state[p2] = copy_clone(opt.state.pop(p1, {}))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "ead99755",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"class ModelToHalf(Callback):\n",
|
|
" \"Use with NonNativeMixedPrecision callback (but it needs to run at the very beginning)\"\n",
|
|
" order=-50\n",
|
|
" def before_fit(self): self.learn.model = convert_network(self.model, dtype=torch.float16)\n",
|
|
" def after_fit (self): self.learn.model = convert_network(self.model, dtype=torch.float32)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "86746ab6",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@docs\n",
|
|
"class NonNativeMixedPrecision(Callback):\n",
|
|
" \"Run training in mixed precision\"\n",
|
|
" order=10\n",
|
|
" def __init__(self, \n",
|
|
" loss_scale:int=512, # Non-dynamic loss scale, used to avoid underflow of gradients. \n",
|
|
" flat_master:bool=False, # Whether to flatten fp32 parameters for performance\n",
|
|
" dynamic:bool=True, # Whether to automatically determine loss scaling\n",
|
|
" max_loss_scale:float=2.**24, # Starting value for dynamic loss scaling\n",
|
|
" div_factor:float=2., # Divide by this on overflow, multiply by this after scale_wait batches\n",
|
|
" scale_wait:int=500, # Number of batches to wait for increasing loss scale\n",
|
|
" clip:float=None, # Value to clip gradients at, max_norm, as in `nn.utils.clip_grad_norm_`\n",
|
|
" ): \n",
|
|
" assert torch.backends.cudnn.enabled, \"Mixed precision training requires cudnn.\"\n",
|
|
" self.flat_master,self.dynamic,self.max_loss_scale = flat_master,dynamic,max_loss_scale\n",
|
|
" self.div_factor,self.scale_wait,self.clip = div_factor,scale_wait,clip\n",
|
|
" self.loss_scale = max_loss_scale if dynamic else loss_scale\n",
|
|
"\n",
|
|
" def before_fit(self):\n",
|
|
" assert self.dls.device.type == 'cuda', \"Mixed-precision training requires a GPU, remove the call `to_fp16`\"\n",
|
|
" if self.learn.opt is None: self.learn.create_opt()\n",
|
|
" self.model_pgs,self.master_pgs = get_master(self.opt, self.flat_master)\n",
|
|
" self.old_pgs = self.opt.param_lists\n",
|
|
" #Changes the optimizer so that the optimization step is done in FP32.\n",
|
|
" _copy_state(self.learn.opt, self.model_pgs, self.master_pgs)\n",
|
|
" if self.dynamic: self.count = 0\n",
|
|
"\n",
|
|
" def before_batch(self): self.learn.xb = to_half(self.xb)\n",
|
|
" def after_pred(self): self.learn.pred = to_float(self.pred)\n",
|
|
" def before_backward(self): self.learn.loss_grad *= self.loss_scale\n",
|
|
"\n",
|
|
" def before_step(self):\n",
|
|
" #First, check for an overflow\n",
|
|
" if self.dynamic and grad_overflow(self.model_pgs):\n",
|
|
" self.loss_scale /= self.div_factor\n",
|
|
" self.learn.loss_grad /= self.div_factor #to record correct loss\n",
|
|
" self.model.zero_grad()\n",
|
|
" raise CancelBatchException() #skip step and zero_grad\n",
|
|
" to_master_grads(self.model_pgs, self.master_pgs, self.flat_master)\n",
|
|
" for master_params in self.master_pgs:\n",
|
|
" for param in master_params:\n",
|
|
" if param.grad is not None: param.grad.div_(self.loss_scale)\n",
|
|
" if self.clip is not None:\n",
|
|
" for group in self.master_pgs: nn.utils.clip_grad_norm_(group, self.clip)\n",
|
|
" # Check if it's been long enough without overflow\n",
|
|
" if self.dynamic:\n",
|
|
" self.count += 1\n",
|
|
" if self.count == self.scale_wait:\n",
|
|
" self.count = 0\n",
|
|
" self.loss_scale *= self.div_factor\n",
|
|
"\n",
|
|
" def after_step(self):\n",
|
|
" self.model.zero_grad() #Zero the gradients of the model manually (optimizer disconnected)\n",
|
|
" to_model_params(self.model_pgs, self.master_pgs, self.flat_master)\n",
|
|
"\n",
|
|
" def after_batch(self):\n",
|
|
" if self.training: self.learn.loss_grad /= self.loss_scale #Log correct loss\n",
|
|
" def after_fit(self):\n",
|
|
" if not hasattr(self,'master_pgs'): return\n",
|
|
" _copy_state(self.learn.opt, self.master_pgs, self.model_pgs)\n",
|
|
" self.learn.opt.param_lists = self.old_pgs\n",
|
|
" delattr(self, \"master_pgs\")\n",
|
|
" delattr(self, \"model_pgs\")\n",
|
|
" delattr(self, \"old_pgs\")\n",
|
|
"\n",
|
|
" _docs = dict(before_fit=\"Put the model in FP16 and prepare the two copies of the parameters\",\n",
|
|
" before_batch=\"Put the input in FP16\",\n",
|
|
" after_pred=\"Put the output back to FP32 so that the loss is computed in FP32\",\n",
|
|
" before_backward=\"Apply loss scaling to avoid gradient underflow\",\n",
|
|
" before_step=\"Update and apply dynamic loss scaling, move gradients to fp32, apply gradient clipping\",\n",
|
|
" after_step=\"Zero fp16 grads and update fp16 params with fp32 params. \",\n",
|
|
" after_batch=\"Ensure loss is logged correctly\",\n",
|
|
" after_fit=\"Put the model back in FP32\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f8fcbfde",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"class TestBeforeMixedPrecision(Callback):\n",
|
|
" order=-55\n",
|
|
" def before_fit(self): test_eq(first(self.model.parameters()).dtype, torch.float32)\n",
|
|
" def before_batch(self): test_eq(self.x.dtype, torch.float32)\n",
|
|
" def after_pred(self): test_eq(self.pred.dtype, torch.float16)\n",
|
|
" def after_loss(self): self.tst_loss = self.learn.loss_grad.detach().clone()\n",
|
|
" def before_step(self):\n",
|
|
" self.learn.has_overflown = grad_overflow(self.non_native_mixed_precision.model_pgs)\n",
|
|
" self.grads = [p.grad.data.clone() for p in self.model.parameters()]\n",
|
|
" self.old_params = [p.data.clone() for p in self.model.parameters()]\n",
|
|
" def after_cancel_step(self): assert self.has_overflown\n",
|
|
"\n",
|
|
"class TestAfterMixedPrecision(Callback):\n",
|
|
" order=65\n",
|
|
" def before_fit(self): test_eq(first(self.model.parameters()).dtype, torch.float16)\n",
|
|
" def after_fit(self): test_eq(first(self.model.parameters()).dtype, torch.float32)\n",
|
|
" def before_batch(self): test_eq(self.x.dtype, torch.float16)\n",
|
|
" def after_pred(self): test_eq(self.pred.dtype, torch.float32)\n",
|
|
" def before_backward(self):\n",
|
|
" loss_scale = self.non_native_mixed_precision.loss_scale if self.training else 1.\n",
|
|
" test_eq(self.loss_grad, self.test_before_mixed_precision.tst_loss * loss_scale) \n",
|
|
" def before_step(self):\n",
|
|
" tbmp = self.test_before_mixed_precision\n",
|
|
" test_eq(self.loss_grad, tbmp.loss_grad)\n",
|
|
" #Test gradients have been copied and scaled back\n",
|
|
" test_close(sum([[p.grad.data for p in pg] for pg in self.non_native_mixed_precision.master_pgs], []),\n",
|
|
" [g.float()/self.non_native_mixed_precision.loss_scale for g in tbmp.grads])\n",
|
|
" def after_batch(self):\n",
|
|
" if self.has_overflown: return\n",
|
|
" tbmp,mp =self.test_before_mixed_precision,self.non_native_mixed_precision\n",
|
|
" #Test master params have been copied to model\n",
|
|
" test_close(sum([[p.data for p in pg] for pg in mp.master_pgs], []),\n",
|
|
" [p.data.float() for p in self.model.parameters()], eps=1e-3)\n",
|
|
" #Test update has been done properly\n",
|
|
" for p,g,op in zip(self.model.parameters(), tbmp.grads, tbmp.old_params):\n",
|
|
" test_close(p.data.float(), op.float() - self.lr*g.float()/self.non_native_mixed_precision.loss_scale, eps=1e-3)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "37738da4",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/html": [
|
|
"<table border=\"1\" class=\"dataframe\">\n",
|
|
" <thead>\n",
|
|
" <tr style=\"text-align: left;\">\n",
|
|
" <th>epoch</th>\n",
|
|
" <th>train_loss</th>\n",
|
|
" <th>valid_loss</th>\n",
|
|
" <th>time</th>\n",
|
|
" </tr>\n",
|
|
" </thead>\n",
|
|
" <tbody>\n",
|
|
" <tr>\n",
|
|
" <td>0</td>\n",
|
|
" <td>7.187932</td>\n",
|
|
" <td>5.855845</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>1</td>\n",
|
|
" <td>7.148743</td>\n",
|
|
" <td>5.697717</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>2</td>\n",
|
|
" <td>7.048915</td>\n",
|
|
" <td>5.524172</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" </tbody>\n",
|
|
"</table>"
|
|
],
|
|
"text/plain": [
|
|
"<IPython.core.display.HTML object>"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
}
|
|
],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| cuda\n",
|
|
"learn = synth_learner(cbs=[ModelToHalf(), NonNativeMixedPrecision()], cuda=True)\n",
|
|
"learn.model = nn.Sequential(nn.Linear(1,1), nn.Linear(1,1)).cuda()\n",
|
|
"learn.opt_func = partial(SGD, mom=0.)\n",
|
|
"learn.splitter = lambda m: [list(m[0].parameters()), list(m[1].parameters())]\n",
|
|
"learn.fit(3, cbs=[TestAfterMixedPrecision(), TestBeforeMixedPrecision()])\n",
|
|
"#Check loss scale did change\n",
|
|
"assert 1 < learn.non_native_mixed_precision.loss_scale < 2**24\n",
|
|
"#Check the model did train\n",
|
|
"for v1,v2 in zip(learn.recorder.values[0], learn.recorder.values[-1]): assert v2<v1"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d3234557",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/html": [
|
|
"<table border=\"1\" class=\"dataframe\">\n",
|
|
" <thead>\n",
|
|
" <tr style=\"text-align: left;\">\n",
|
|
" <th>epoch</th>\n",
|
|
" <th>train_loss</th>\n",
|
|
" <th>valid_loss</th>\n",
|
|
" <th>time</th>\n",
|
|
" </tr>\n",
|
|
" </thead>\n",
|
|
" <tbody>\n",
|
|
" <tr>\n",
|
|
" <td>0</td>\n",
|
|
" <td>11.927933</td>\n",
|
|
" <td>12.063744</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>1</td>\n",
|
|
" <td>11.539829</td>\n",
|
|
" <td>11.545557</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>2</td>\n",
|
|
" <td>11.266481</td>\n",
|
|
" <td>11.075830</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" </tbody>\n",
|
|
"</table>"
|
|
],
|
|
"text/plain": [
|
|
"<IPython.core.display.HTML object>"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
}
|
|
],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| cuda\n",
|
|
"learn = synth_learner(cbs=[ModelToHalf(), NonNativeMixedPrecision(dynamic=False)], cuda=True)\n",
|
|
"learn.model = nn.Sequential(nn.Linear(1,1), nn.Linear(1,1)).cuda()\n",
|
|
"learn.opt_func = partial(SGD, mom=0.)\n",
|
|
"learn.splitter = lambda m: [list(m[0].parameters()), list(m[1].parameters())]\n",
|
|
"learn.fit(3, cbs=[TestAfterMixedPrecision(), TestBeforeMixedPrecision()])\n",
|
|
"#Check loss scale did mot change\n",
|
|
"test_eq(learn.non_native_mixed_precision.loss_scale,512)\n",
|
|
"#Check the model did train\n",
|
|
"for v1,v2 in zip(learn.recorder.values[0], learn.recorder.values[-1]): assert v2<v1"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "571f71c1",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"@delegates(NonNativeMixedPrecision.__init__)\n",
|
|
"def to_non_native_fp16(self:Learner, **kwargs): return self.add_cbs([ModelToHalf(), NonNativeMixedPrecision(**kwargs)])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9ec2796c",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/html": [
|
|
"<table border=\"1\" class=\"dataframe\">\n",
|
|
" <thead>\n",
|
|
" <tr style=\"text-align: left;\">\n",
|
|
" <th>epoch</th>\n",
|
|
" <th>train_loss</th>\n",
|
|
" <th>valid_loss</th>\n",
|
|
" <th>time</th>\n",
|
|
" </tr>\n",
|
|
" </thead>\n",
|
|
" <tbody>\n",
|
|
" <tr>\n",
|
|
" <td>0</td>\n",
|
|
" <td>8.358611</td>\n",
|
|
" <td>10.943352</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>1</td>\n",
|
|
" <td>8.330508</td>\n",
|
|
" <td>10.722443</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>2</td>\n",
|
|
" <td>8.221409</td>\n",
|
|
" <td>10.485508</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" </tbody>\n",
|
|
"</table>"
|
|
],
|
|
"text/plain": [
|
|
"<IPython.core.display.HTML object>"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
}
|
|
],
|
|
"source": [
|
|
"#| cuda\n",
|
|
"learn = synth_learner(cuda=True)\n",
|
|
"learn.model = nn.Sequential(nn.Linear(1,1), nn.Linear(1,1)).cuda()\n",
|
|
"learn.opt_func = partial(SGD, mom=0.)\n",
|
|
"learn.splitter = lambda m: [list(m[0].parameters()), list(m[1].parameters())]\n",
|
|
"learn.to_non_native_fp16()\n",
|
|
"learn.fit(3, cbs=[TestAfterMixedPrecision(), TestBeforeMixedPrecision()])\n",
|
|
"#Check the model did train\n",
|
|
"for v1,v2 in zip(learn.recorder.values[0], learn.recorder.values[-1]): assert v2<v1"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9d76a9af",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/html": [
|
|
"<table border=\"1\" class=\"dataframe\">\n",
|
|
" <thead>\n",
|
|
" <tr style=\"text-align: left;\">\n",
|
|
" <th>epoch</th>\n",
|
|
" <th>train_loss</th>\n",
|
|
" <th>valid_loss</th>\n",
|
|
" <th>time</th>\n",
|
|
" </tr>\n",
|
|
" </thead>\n",
|
|
" <tbody>\n",
|
|
" <tr>\n",
|
|
" <td>0</td>\n",
|
|
" <td>11.646567</td>\n",
|
|
" <td>10.883919</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>1</td>\n",
|
|
" <td>11.489956</td>\n",
|
|
" <td>9.904404</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>2</td>\n",
|
|
" <td>10.746455</td>\n",
|
|
" <td>7.914827</td>\n",
|
|
" <td>00:00</td>\n",
|
|
" </tr>\n",
|
|
" </tbody>\n",
|
|
"</table>"
|
|
],
|
|
"text/plain": [
|
|
"<IPython.core.display.HTML object>"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
}
|
|
],
|
|
"source": [
|
|
"#| hide\n",
|
|
"#| cuda\n",
|
|
"learn = synth_learner(cuda=True)\n",
|
|
"learn.model = nn.Sequential(nn.Linear(1,1), nn.Linear(1,1)).cuda()\n",
|
|
"learn.opt_func = partial(SGD, mom=0.9)\n",
|
|
"learn.splitter = lambda m: [list(m[0].parameters()), list(m[1].parameters())]\n",
|
|
"learn.to_non_native_fp16()\n",
|
|
"learn.freeze()\n",
|
|
"learn.create_opt()\n",
|
|
"init_ps = [p for pg in learn.opt.param_groups for p in pg]\n",
|
|
"learn.fit(3)\n",
|
|
"final_ps = [p for pg in learn.opt.param_groups for p in pg]\n",
|
|
"for p1,p2 in zip(init_ps, final_ps): test_is(p1, p2)\n",
|
|
"#First param groups has no state because not trained\n",
|
|
"test_eq([learn.opt.state[p] for p in learn.opt.param_lists[0]], [{}, {'do_wd': False}])\n",
|
|
"#Second param groups has state \n",
|
|
"for p in learn.opt.param_lists[1]: assert 'grad_avg' in learn.opt.state[p]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f202da47",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| export\n",
|
|
"@patch\n",
|
|
"def to_non_native_fp32(self: Learner): return self.remove_cbs([ModelToHalf, NonNativeMixedPrecision])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1821affb",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| cuda\n",
|
|
"learn = learn.to_non_native_fp32()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "ab083f25",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Export -"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9d60eca8",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"#| hide\n",
|
|
"from nbdev import *\n",
|
|
"nbdev_export()"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"jupytext": {
|
|
"split_at_heading": true
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "python3",
|
|
"language": "python",
|
|
"name": "python3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|