{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "## Hook callbacks" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "This provides both a standalone class and a callback for registering and automatically deregistering [PyTorch hooks](https://pytorch.org/tutorials/beginner/former_torchies/nn_tutorial.html#forward-and-backward-function-hooks), along with some pre-defined hooks. Hooks can be attached to any [`nn.Module`](https://pytorch.org/docs/stable/nn.html#torch.nn.Module), for either the forward or the backward pass.\n", "\n", "We'll start by looking at the pre-defined hook [`ActivationStats`](/callbacks.hooks.html#ActivationStats), then we'll see how to create our own." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "hide_input": true }, "outputs": [], "source": [ "from fastai.gen_doc.nbdoc import *\n", "from fastai.callbacks.hooks import * \n", "from fastai.train import *\n", "from fastai.vision import *" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "hide_input": true }, "outputs": [ { "data": { "text/markdown": [ "
class ActivationStats[source]ActivationStats(**`learn`**:[`Learner`](/basic_train.html#Learner), **`modules`**:`Sequence`\\[[`Module`](https://pytorch.org/docs/stable/nn.html#torch.nn.Module)\\]=***`None`***, **`do_remove`**:`bool`=***`True`***) :: [`HookCallback`](/callbacks.hooks.html#HookCallback)\n",
"\n",
"Callback that record the mean and std of activations. "
],
"text/plain": [
"| epoch | \n", "train_loss | \n", "valid_loss | \n", "
|---|---|---|
| 1 | \n", "0.112384 | \n", "0.083544 | \n", "
hook[source]hook(**`m`**:[`Module`](https://pytorch.org/docs/stable/nn.html#torch.nn.Module), **`i`**:`Tensors`, **`o`**:`Tensors`) → `Tuple`\\[`Rank0Tensor`, `Rank0Tensor`\\]\n",
"\n",
"Take the mean and std of `o`. "
],
"text/plain": [
"on_train_begin[source]on_train_begin(**\\*\\*`kwargs`**)\n",
"\n",
"Initialize stats. "
],
"text/plain": [
"on_batch_end[source]on_batch_end(**`train`**, **\\*\\*`kwargs`**)\n",
"\n",
"Take the stored results and puts it in `self.stats` "
],
"text/plain": [
"on_train_end[source]on_train_end(**\\*\\*`kwargs`**)\n",
"\n",
"Polish the final result. "
],
"text/plain": [
"class Hook[source]Hook(**`m`**:[`Module`](https://pytorch.org/docs/stable/nn.html#torch.nn.Module), **`hook_func`**:`HookFunc`, **`is_forward`**:`bool`=***`True`***, **`detach`**:`bool`=***`True`***)\n",
"\n",
"Create a hook on `m` with `hook_func`. "
],
"text/plain": [
"remove[source]remove()\n",
"\n",
"Remove the hook from the model. "
],
"text/plain": [
"class Hooks[source]Hooks(**`ms`**:`ModuleList`, **`hook_func`**:`HookFunc`, **`is_forward`**:`bool`=***`True`***, **`detach`**:`bool`=***`True`***)\n",
"\n",
"Create several hooks on the modules in `ms` with `hook_func`. "
],
"text/plain": [
"remove[source]remove()\n",
"\n",
"Remove the hooks from the model. "
],
"text/plain": [
"hook_output[source]hook_output(**`module`**:[`Module`](https://pytorch.org/docs/stable/nn.html#torch.nn.Module), **`detach`**:`bool`=***`True`***, **`grad`**:`bool`=***`False`***) → [`Hook`](/callbacks.hooks.html#Hook)\n",
"\n",
"Return a [`Hook`](/callbacks.hooks.html#Hook) that stores activations of `module` in `self.stored` "
],
"text/plain": [
"hook_outputs[source]hook_outputs(**`modules`**:`ModuleList`, **`detach`**:`bool`=***`True`***, **`grad`**:`bool`=***`False`***) → [`Hooks`](/callbacks.hooks.html#Hooks)\n",
"\n",
"Return [`Hooks`](/callbacks.hooks.html#Hooks) that store activations of all `modules` in `self.stored` "
],
"text/plain": [
"model_sizes[source]model_sizes(**`m`**:[`Module`](https://pytorch.org/docs/stable/nn.html#torch.nn.Module), **`size`**:`tuple`=***`(64, 64)`***) → `Tuple`\\[`Sizes`, `Tensor`, [`Hooks`](/callbacks.hooks.html#Hooks)\\]\n",
"\n",
"Pass a dummy input through the model `m` to get the various sizes of activations. "
],
"text/plain": [
"model_summary[source]model_summary(**`m`**:[`Learner`](/basic_train.html#Learner), **`n`**:`int`=***`70`***)\n",
"\n",
"Print a summary of `m` using a output text width of `n` chars "
],
"text/plain": [
"num_features_model[source]num_features_model(**`m`**:[`Module`](https://pytorch.org/docs/stable/nn.html#torch.nn.Module)) → `int`\n",
"\n",
"Return the number of output features for `model`. "
],
"text/plain": [
"dummy_batch[source]dummy_batch(**`m`**:[`Module`](https://pytorch.org/docs/stable/nn.html#torch.nn.Module), **`size`**:`tuple`=***`(64, 64)`***) → `Tensor`\n",
"\n",
"Create a dummy batch to go through `m` with `size`. "
],
"text/plain": [
"dummy_eval[source]dummy_eval(**`m`**:[`Module`](https://pytorch.org/docs/stable/nn.html#torch.nn.Module), **`size`**:`tuple`=***`(64, 64)`***)\n",
"\n",
"Pass a [`dummy_batch`](/callbacks.hooks.html#dummy_batch) in evaluation mode in `m` with `size`. "
],
"text/plain": [
"class HookCallback[source]HookCallback(**`learn`**:[`Learner`](/basic_train.html#Learner), **`modules`**:`Sequence`\\[[`Module`](https://pytorch.org/docs/stable/nn.html#torch.nn.Module)\\]=***`None`***, **`do_remove`**:`bool`=***`True`***) :: [`LearnerCallback`](/basic_train.html#LearnerCallback)\n",
"\n",
"Callback that can be used to register hooks on `modules`. Implement the corresponding function in `self.hook`. "
],
"text/plain": [
"on_train_begin[source]on_train_begin(**\\*\\*`kwargs`**)\n",
"\n",
"Register the [`Hooks`](/callbacks.hooks.html#Hooks) on `self.modules`. "
],
"text/plain": [
"on_train_end[source]on_train_end(**\\*\\*`kwargs`**)\n",
"\n",
"Remove the [`Hooks`](/callbacks.hooks.html#Hooks). "
],
"text/plain": [
"remove[source]remove()"
],
"text/plain": [
"hook_fn[source]hook_fn(**`module`**:[`Module`](https://pytorch.org/docs/stable/nn.html#torch.nn.Module), **`input`**:`Tensors`, **`output`**:`Tensors`)\n",
"\n",
"Applies `hook_func` to `module`, `input`, `output`. "
],
"text/plain": [
"