Skip to content

Instantly share code, notes, and snippets.

@kunalb

kunalb/Torch Compile + Hooks.ipynb

Created January 3, 2023 22:37
Show Gist options
  • Save kunalb/6c3eae96636f972d73b6357f92935980 to your computer and use it in GitHub Desktop.
Save kunalb/6c3eae96636f972d73b6357f92935980 to your computer and use it in GitHub Desktop.
Download ZIP
Display the source blob
Display the rendered blob
Raw
Torch Compile + Hooks.ipynb
{
"metadata": {
"dataExplorerConfig": {},
"bento_stylesheets": {
"bento/extensions/flow/main.css": true,
"bento/extensions/kernel_selector/main.css": true,
"bento/extensions/kernel_ui/main.css": true,
"bento/extensions/new_kernel/main.css": true,
"bento/extensions/system_usage/main.css": true,
"bento/extensions/theme/main.css": true
},
"kernelspec": {
"display_name": "pytorch",
"language": "python",
"name": "bento_kernel_pytorch",
"cinder_runtime": false,
"ipyflow_runtime": false,
"metadata": {
"kernel_name": "bento_kernel_pytorch",
"nightly_builds": true,
"fbpkg_supported": true,
"cinder_runtime": false,
"ipyflow_runtime": false,
"is_prebuilt": true
}
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3"
},
"last_server_session_id": "678f3e8f-5b00-48d7-9975-4fd503f11dab",
"last_kernel_id": "a1635b17-a06a-47b9-be37-7eca2d24c178",
"last_base_url": "https://devvm179.rva0.facebook.com:8090/",
"last_msg_id": "f66e98f0-9ef07ded991c7ad51ba75f9b_761",
"outputWidgetContext": {}
},
"nbformat": 4,
"nbformat_minor": 2,
"cells": [
{
"cell_type": "code",
"metadata": {
"collapsed": false,
"originalKey": "335625b5-ab9b-428e-a893-8728e4633c9a",
"requestMsgId": "5a8ca9f2-aacc-42b1-805f-0a95441e469f",
"executionStartTime": 1672785127434,
"executionStopTime": 1672785127445
},
"source": [
"import torch\n",
"\n",
"from typing import Tuple"
],
"execution_count": 59,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"originalKey": "e7e7beb3-48d6-4906-bfd2-ddaab69f3d13",
"showInput": true,
"customInput": null,
"collapsed": false,
"requestMsgId": "ffa4458c-01d2-44c4-8982-73b7e434ace4",
"executionStartTime": 1672785127534,
"executionStopTime": 1672785127547
},
"source": [
"class TestModule(torch.nn.Module):\n",
"\n",
" def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
" return 2 * x"
],
"execution_count": 60,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"originalKey": "37d5a847-3fb6-4a1d-9d68-7154e6e09a7d",
"showInput": true,
"customInput": null,
"collapsed": false,
"requestMsgId": "bfedbe60-ef35-451f-92d1-df55f3691496",
"executionStartTime": 1672785127647,
"executionStopTime": 1672785127664
},
"source": [
"m = TestModule()"
],
"execution_count": 61,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"originalKey": "cb7a3971-1c51-4b3b-bfce-fc266010dc17",
"showInput": true,
"customInput": null,
"collapsed": false,
"requestMsgId": "95b6e183-40d0-4dd8-a9c1-4460421dd05c",
"executionStartTime": 1672785127746,
"executionStopTime": 1672785127755
},
"source": [
"def forward_hook(module: torch.nn.Module, inputs: Tuple[torch.Tensor], output: torch.Tensor) -> torch.Tensor:\n",
" print(f\"Called forward hook: {inputs=} {output=}\")\n",
" return 2 * output\n",
"\n",
"def forward_pre_hook(module: torch.nn.Module, inputs: Tuple[torch.Tensor]) -> None:\n",
" print(f\"Called forward pre hook: {inputs=}\")\n",
"\n",
"def backward_hook(module: torch.nn.Module, grad_input: Tuple[torch.Tensor], grad_output: Tuple[torch.Tensor]) -> None:\n",
" print(f\"Called backward hook: {grad_input=} {grad_output=}\")"
],
"execution_count": 62,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"originalKey": "ae3cd953-fc4d-4836-a636-0c714f054bd1",
"showInput": true,
"customInput": null,
"collapsed": false,
"requestMsgId": "1baabf4f-0650-4a7d-915d-b9f42ab076e4",
"executionStartTime": 1672785127847,
"executionStopTime": 1672785127896
},
"source": [
"m.register_forward_hook(forward_hook)\n",
"m.register_forward_pre_hook(pre_forward_hook)\n",
"m.register_full_backward_hook(backward_hook)"
],
"execution_count": 63,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": "<torch.utils.hooks.RemovableHandle at 0x7f1961ca1730>"
},
"metadata": {
"bento_obj_id": "139746991544112"
},
"execution_count": 63
}
]
},
{
"cell_type": "code",
"metadata": {
"originalKey": "02519d8a-89c3-4345-934e-8d8ede44b1e9",
"showInput": true,
"customInput": null,
"collapsed": false,
"requestMsgId": "bbdf7fde-0574-40bd-a052-a01e05d5c730",
"executionStartTime": 1672785127968,
"executionStopTime": 1672785128000
},
"source": [
"t = torch.tensor(1.0, requires_grad=True)\n",
"l = 2 * m(t)\n",
"l"
],
"execution_count": 64,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Called pre-forward hook: inputs=(tensor(1., requires_grad=True),)\nCalled forward hook: inputs=(tensor(1., grad_fn=<BackwardHookFunctionBackward>),) output=tensor(2., grad_fn=<MulBackward0>)\n"
]
},
{
"output_type": "execute_result",
"data": {
"text/plain": "tensor(8., grad_fn=<MulBackward0>)"
},
"metadata": {
"bento_obj_id": "139746991647936"
},
"execution_count": 64
}
]
},
{
"cell_type": "code",
"metadata": {
"originalKey": "00142d31-c85d-4c13-9410-41f206f48d0f",
"showInput": true,
"customInput": null,
"collapsed": false,
"requestMsgId": "85c0658b-516c-49e1-a1c2-af739a7362b6",
"executionStartTime": 1672785128082,
"executionStopTime": 1672785128119
},
"source": [
"l.backward()"
],
"execution_count": 65,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Called backward hook: grad_input=(tensor(8.),) grad_output=(tensor(2.),)\n"
]
}
]
},
{
"cell_type": "code",
"metadata": {
"originalKey": "bf4c031b-734a-4bfd-89d9-e065495fbde2",
"showInput": true,
"customInput": null,
"collapsed": false,
"requestMsgId": "08ddc2b1-d219-4731-b1d7-f6ec284b597c",
"executionStartTime": 1672785128212,
"executionStopTime": 1672785128230
},
"source": [
"cm = torch.compile(m)"
],
"execution_count": 66,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"originalKey": "d1ffff4e-4929-4bbe-854b-347039b8196b",
"showInput": true,
"customInput": null,
"collapsed": false,
"requestMsgId": "620b95f6-a105-4504-b4d2-265211d861d1",
"executionStartTime": 1672785128314,
"executionStopTime": 1672785128350
},
"source": [
"cl = 2 * cm(t)\n",
"cl"
],
"execution_count": 67,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": "tensor(4., grad_fn=<MulBackward0>)"
},
"metadata": {
"bento_obj_id": "139746991171760"
},
"execution_count": 67
}
]
},
{
"cell_type": "code",
"metadata": {
"originalKey": "f90f7d27-b24c-42ba-860c-3408b12b04a4",
"showInput": true,
"customInput": null,
"collapsed": false,
"requestMsgId": "03fdd72a-0827-4f8f-9388-3f73b1b64eef",
"executionStartTime": 1672785128429,
"executionStopTime": 1672785128452
},
"source": [
"cl.backward()"
],
"execution_count": 68,
"outputs": []
}
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment