Skip to content

Commit 83ae20a

Browse files
JKSenthilfacebook-github-bot
JKSenthil
authored and
facebook-github-bot
committed
move ModuleSummary + FLOPS tools from eval to tnt (#169)
Summary: X-link: pytorch/tnt#476 Pull Request resolved: #169 Move `flops`, `ModuleSummary` utilities to TNT, and removed reference to flops in quickstart jupyter notebook Also, module.compile is only in pytorch nightlies, and it seems the nightly pytorch is being overrided by the stable torch package in requirements.txt, so filtering out torch in the github workflows Reviewed By: ananthsub Differential Revision: D47735567 fbshipit-source-id: 8ecabb70cbddeaf27bbd5141bf9ebcafff5eadd4
1 parent e63b8de commit 83ae20a

File tree

6 files changed

+4
-1814
lines changed

6 files changed

+4
-1814
lines changed

‎examples/Introducing_TorchEval.ipynb

+4-156
Original file line numberDiff line numberDiff line change
@@ -604,159 +604,6 @@
604604
"Notice that our final result is computed with 20,000 samples, 5,000 from each process!"
605605
]
606606
},
607-
{
608-
"cell_type": "markdown",
609-
"metadata": {
610-
"id": "6XIhytwSeNPd"
611-
},
612-
"source": [
613-
"# Module Summary Tools\n",
614-
"\n",
615-
"TorchEval also includes tools for model summarization, providing per layer details of trainable parameters, size in bytes, FLOPS, and more."
616-
]
617-
},
618-
{
619-
"cell_type": "markdown",
620-
"metadata": {
621-
"id": "uC9eKXXkg5v0"
622-
},
623-
"source": [
624-
"To get a basic summary of the model, we can use `get_module_summary` and pass in the model\n",
625-
"\n",
626-
"\n",
627-
"\n"
628-
]
629-
},
630-
{
631-
"cell_type": "code",
632-
"execution_count": 14,
633-
"metadata": {
634-
"colab": {
635-
"base_uri": "https://localhost:8080/"
636-
},
637-
"id": "qaZmg2YWe7PS",
638-
"outputId": "50915417-cee2-4e3f-a7d8-2c9b628bc75e"
639-
},
640-
"outputs": [
641-
{
642-
"name": "stdout",
643-
"output_type": "stream",
644-
"text": [
645-
"Name | Type | # Parameters | # Trainable Parameters | Size (bytes) | Contains Uninitialized Parameters?\n",
646-
"----------------------------------------------------------------------------------------------------------------------------\n",
647-
" | AlexNet | 61.1 M | 61.1 M | 244 M | No \n",
648-
"features | Sequential | 2.5 M | 2.5 M | 9.9 M | No \n",
649-
"features.0 | Conv2d | 23.3 K | 23.3 K | 93.2 K | No \n",
650-
"features.1 | ReLU | 0 | 0 | 0 | No \n",
651-
"features.2 | MaxPool2d | 0 | 0 | 0 | No \n",
652-
"features.3 | Conv2d | 307 K | 307 K | 1.2 M | No \n",
653-
"features.4 | ReLU | 0 | 0 | 0 | No \n",
654-
"features.5 | MaxPool2d | 0 | 0 | 0 | No \n",
655-
"features.6 | Conv2d | 663 K | 663 K | 2.7 M | No \n",
656-
"features.7 | ReLU | 0 | 0 | 0 | No \n",
657-
"features.8 | Conv2d | 884 K | 884 K | 3.5 M | No \n",
658-
"features.9 | ReLU | 0 | 0 | 0 | No \n",
659-
"features.10 | Conv2d | 590 K | 590 K | 2.4 M | No \n",
660-
"features.11 | ReLU | 0 | 0 | 0 | No \n",
661-
"features.12 | MaxPool2d | 0 | 0 | 0 | No \n",
662-
"avgpool | AdaptiveAvgPool2d | 0 | 0 | 0 | No \n",
663-
"classifier | Sequential | 58.6 M | 58.6 M | 234 M | No \n",
664-
"classifier.0 | Dropout | 0 | 0 | 0 | No \n",
665-
"classifier.1 | Linear | 37.8 M | 37.8 M | 151 M | No \n",
666-
"classifier.2 | ReLU | 0 | 0 | 0 | No \n",
667-
"classifier.3 | Dropout | 0 | 0 | 0 | No \n",
668-
"classifier.4 | Linear | 16.8 M | 16.8 M | 67.1 M | No \n",
669-
"classifier.5 | ReLU | 0 | 0 | 0 | No \n",
670-
"classifier.6 | Linear | 4.1 M | 4.1 M | 16.4 M | No \n",
671-
"\n"
672-
]
673-
}
674-
],
675-
"source": [
676-
"from torchvision.models.alexnet import AlexNet\n",
677-
"from torcheval.tools import get_module_summary\n",
678-
"\n",
679-
"model = AlexNet()\n",
680-
"ms = get_module_summary(model)\n",
681-
"print(ms)"
682-
]
683-
},
684-
{
685-
"cell_type": "markdown",
686-
"metadata": {},
687-
"source": [
688-
"In the table above we see the layers of AlexNet printed, alongside the parameter count and size in bytes at each layer."
689-
]
690-
},
691-
{
692-
"cell_type": "markdown",
693-
"metadata": {
694-
"id": "8gx5fSgVf1C1"
695-
},
696-
"source": [
697-
"Passing in an example input tensor will retrieve additional metrics such as FLOPS (number of multiply-add operations), activation sizes, and more"
698-
]
699-
},
700-
{
701-
"cell_type": "code",
702-
"execution_count": 15,
703-
"metadata": {
704-
"colab": {
705-
"base_uri": "https://localhost:8080/"
706-
},
707-
"id": "4UdzYuflf0Iz",
708-
"outputId": "da16fd14-f8c8-40dd-eb27-105bef69c21c"
709-
},
710-
"outputs": [
711-
{
712-
"name": "stdout",
713-
"output_type": "stream",
714-
"text": [
715-
"Name | Type | # Parameters | # Trainable Parameters | Size (bytes) | Contains Uninitialized Parameters? | Forward FLOPs | Backward FLOPs | In size | Out size | Forward Elapsed Times (ms)\n",
716-
"--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\n",
717-
" | AlexNet | 61.1 M | 61.1 M | 244 M | No | 714 M | 1.4 G | [1, 3, 224, 224] | [1, 1000] | 0.0004032743 \n",
718-
"features | Sequential | 2.5 M | 2.5 M | 9.9 M | No | 655 M | 1.2 G | [1, 3, 224, 224] | [1, 256, 6, 6] | 0.0003372050 \n",
719-
"features.0 | Conv2d | 23.3 K | 23.3 K | 93.2 K | No | 70.3 M | 70.3 M | [1, 3, 224, 224] | [1, 64, 55, 55] | 0.0002422328 \n",
720-
"features.1 | ReLU | 0 | 0 | 0 | No | 0 | 0 | [1, 64, 55, 55] | [1, 64, 55, 55] | 0.0000035029 \n",
721-
"features.2 | MaxPool2d | 0 | 0 | 0 | No | 0 | 0 | [1, 64, 55, 55] | [1, 64, 27, 27] | 0.0000089833 \n",
722-
"features.3 | Conv2d | 307 K | 307 K | 1.2 M | No | 223 M | 447 M | [1, 64, 27, 27] | [1, 192, 27, 27] | 0.0000217149 \n",
723-
"features.4 | ReLU | 0 | 0 | 0 | No | 0 | 0 | [1, 192, 27, 27] | [1, 192, 27, 27] | 0.0000007634 \n",
724-
"features.5 | MaxPool2d | 0 | 0 | 0 | No | 0 | 0 | [1, 192, 27, 27] | [1, 192, 13, 13] | 0.0000033676 \n",
725-
"features.6 | Conv2d | 663 K | 663 K | 2.7 M | No | 112 M | 224 M | [1, 192, 13, 13] | [1, 384, 13, 13] | 0.0000095535 \n",
726-
"features.7 | ReLU | 0 | 0 | 0 | No | 0 | 0 | [1, 384, 13, 13] | [1, 384, 13, 13] | 0.0000006440 \n",
727-
"features.8 | Conv2d | 884 K | 884 K | 3.5 M | No | 149 M | 299 M | [1, 384, 13, 13] | [1, 256, 13, 13] | 0.0000138996 \n",
728-
"features.9 | ReLU | 0 | 0 | 0 | No | 0 | 0 | [1, 256, 13, 13] | [1, 256, 13, 13] | 0.0000005126 \n",
729-
"features.10 | Conv2d | 590 K | 590 K | 2.4 M | No | 99.7 M | 199 M | [1, 256, 13, 13] | [1, 256, 13, 13] | 0.0000107368 \n",
730-
"features.11 | ReLU | 0 | 0 | 0 | No | 0 | 0 | [1, 256, 13, 13] | [1, 256, 13, 13] | 0.0000007307 \n",
731-
"features.12 | MaxPool2d | 0 | 0 | 0 | No | 0 | 0 | [1, 256, 13, 13] | [1, 256, 6, 6] | 0.0000012592 \n",
732-
"avgpool | AdaptiveAvgPool2d | 0 | 0 | 0 | No | 0 | 0 | [1, 256, 6, 6] | [1, 256, 6, 6] | 0.0000035178 \n",
733-
"classifier | Sequential | 58.6 M | 58.6 M | 234 M | No | 58.6 M | 117 M | [1, 9216] | [1, 1000] | 0.0000574685 \n",
734-
"classifier.0 | Dropout | 0 | 0 | 0 | No | 0 | 0 | [1, 9216] | [1, 9216] | 0.0000154559 \n",
735-
"classifier.1 | Linear | 37.8 M | 37.8 M | 151 M | No | 37.7 M | 75.5 M | [1, 9216] | [1, 4096] | 0.0000240998 \n",
736-
"classifier.2 | ReLU | 0 | 0 | 0 | No | 0 | 0 | [1, 4096] | [1, 4096] | 0.0000004282 \n",
737-
"classifier.3 | Dropout | 0 | 0 | 0 | No | 0 | 0 | [1, 4096] | [1, 4096] | 0.0000006650 \n",
738-
"classifier.4 | Linear | 16.8 M | 16.8 M | 67.1 M | No | 16.8 M | 33.6 M | [1, 4096] | [1, 4096] | 0.0000103166 \n",
739-
"classifier.5 | ReLU | 0 | 0 | 0 | No | 0 | 0 | [1, 4096] | [1, 4096] | 0.0000010071 \n",
740-
"classifier.6 | Linear | 4.1 M | 4.1 M | 16.4 M | No | 4.1 M | 8.2 M | [1, 4096] | [1, 1000] | 0.0000022619 \n",
741-
"Remark for FLOPs calculation: (1) Only operators `mm`|`matmul`|`addmm`|`bmm`|`convolution`|`_convolution`|`convolution_backward` are included. To add more operators supported in FLOPs calculation, please contribute to torcheval/tools/flops.py. (2) The calculation related to additional loss function is not included. For forward, we calculated FLOPs based on `loss = model(input_data).mean()`. For backward, we calculated FLOPs based on `loss.backward()`. \n",
742-
"\n"
743-
]
744-
}
745-
],
746-
"source": [
747-
"model = AlexNet()\n",
748-
"inp = torch.randn(1, 3, 224, 224)\n",
749-
"ms = get_module_summary(model, module_args=(inp,))\n",
750-
"print(ms)"
751-
]
752-
},
753-
{
754-
"cell_type": "markdown",
755-
"metadata": {},
756-
"source": [
757-
"By passing this tensor through the model, the module summary is additionally able to provide FLOPS, activations sizes, and time elapsed at each layer. FLOPS are computed by utilizing [TorchDispatchMode](https://dev-discuss.pytorch.org/t/torchdispatchmode-for-debugging-testing-and-more/717) to interpose at the [__torch_dispatch__](https://dev-discuss.pytorch.org/t/what-and-why-is-torch-dispatch/557) level, where all operators on the input tensor are caught and FLOPS at each operator are computed and added together."
758-
]
759-
},
760607
{
761608
"cell_type": "markdown",
762609
"metadata": {
@@ -775,10 +622,11 @@
775622
"colab": {
776623
"provenance": []
777624
},
625+
"fileHeader": "",
778626
"kernelspec": {
779-
"display_name": "Python 3.10.6 ('evaldocs')",
627+
"display_name": "Python 3",
780628
"language": "python",
781-
"name": "python3"
629+
"name": "bento_kernel_default"
782630
},
783631
"language_info": {
784632
"codemirror_mode": {
@@ -799,5 +647,5 @@
799647
}
800648
},
801649
"nbformat": 4,
802-
"nbformat_minor": 0
650+
"nbformat_minor": 2
803651
}

0 commit comments

Comments
 (0)