Add flop counter utility #95751

Chillee · 2023-03-01T01:24:39Z

Stack from ghstack (oldest at bottom):

Overall, an example usage. Note that this also captures backwards FLOPs.

import torchvision.models as models
import torch
from torch.utils.flop_counter import FlopCounterMode

inp = torch.randn(1, 3, 224, 224, device='cpu')
mod = models.resnet18()

flop_counter = FlopCounterMode(mod, depth=1)
with flop_counter:
    mod(inp).sum().backward()

You can control the depth of the module hierarchy with the depth attribute (which defaults to 2). For example, if I don't limit it, this is what it outputs.

Other APIs

FlopCounterMode(custom_mapping=...): Allows for custom flop counting functions
FlopCounterMode.get_table(depth=...): Explicitly get the table as a string
FlopCounterMode.flop_counts: Contains the flop information as a Dict[hierarchy: str, Dict[Op, int]]
FlopCounterMode.register_hierarchy(f, name): Allows you to register additional "hierarchies" for a function.

[ghstack-poisoned]

pytorch-bot · 2023-03-01T01:24:42Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/95751

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit d7256a4:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

ghstack-source-id: 477bdd302c406e1c4c7e85ab9eb67dc2c9ce5da2 Pull Request resolved: #95751

[ghstack-poisoned]

ghstack-source-id: 871593381439193148689ea823e158338cbfa154 Pull Request resolved: #95751

[ghstack-poisoned]

ghstack-source-id: a43cde0a4933ebd1a1839f57c2840d7da7b94d39 Pull Request resolved: #95751

Overall, an example usage. Note that this *also* captures backwards FLOPs. ``` import torchvision.models as models import torch from torch.utils.flop_counter import FlopCounterMode inp = torch.randn(1, 3, 224, 224, device='cpu') mod = models.resnet18() flop_counter = FlopCounterMode(mod, depth=1) with flop_counter: mod(inp).sum().backward() ``` <img width="326" alt="image" src="https://user-images.githubusercontent.com/6355099/222023068-3491e405-f195-4e11-b679-36b19a1380c7.png"> You can control the depth of the module hierarchy with the `depth` attribute (which defaults to 2). For example, if I don't limit it, this is what it outputs. <img width="366" alt="image" src="https://user-images.githubusercontent.com/6355099/222023306-3d880bb6-f534-4f98-bf10-83c4353acefc.png"> ## Other APIs FlopCounterMode(custom_mapping=...): Allows for custom flop counting functions FlopCounterMode.get_table(depth=...): Explicitly get the table as a string FlopCounterMode.flop_counts: Contains the flop information as a Dict[hierarchy: str, Dict[Op, int]] [ghstack-poisoned]

ghstack-source-id: 59eb9a3290d67db39822b4d14d163e2b29f811cd Pull Request resolved: #95751

Overall, an example usage. Note that this *also* captures backwards FLOPs. ``` import torchvision.models as models import torch from torch.utils.flop_counter import FlopCounterMode inp = torch.randn(1, 3, 224, 224, device='cpu') mod = models.resnet18() flop_counter = FlopCounterMode(mod, depth=1) with flop_counter: mod(inp).sum().backward() ``` <img width="326" alt="image" src="https://user-images.githubusercontent.com/6355099/222023068-3491e405-f195-4e11-b679-36b19a1380c7.png"> You can control the depth of the module hierarchy with the `depth` attribute (which defaults to 2). For example, if I don't limit it, this is what it outputs. <img width="366" alt="image" src="https://user-images.githubusercontent.com/6355099/222023306-3d880bb6-f534-4f98-bf10-83c4353acefc.png"> ## Other APIs FlopCounterMode(custom_mapping=...): Allows for custom flop counting functions FlopCounterMode.get_table(depth=...): Explicitly get the table as a string FlopCounterMode.flop_counts: Contains the flop information as a Dict[hierarchy: str, Dict[Op, int]] [ghstack-poisoned]

ghstack-source-id: eb8901eff886f9018a19fb5d5730d1139069064b Pull Request resolved: #95751

Overall, an example usage. Note that this *also* captures backwards FLOPs. ``` import torchvision.models as models import torch from torch.utils.flop_counter import FlopCounterMode inp = torch.randn(1, 3, 224, 224, device='cpu') mod = models.resnet18() flop_counter = FlopCounterMode(mod, depth=1) with flop_counter: mod(inp).sum().backward() ``` <img width="326" alt="image" src="https://user-images.githubusercontent.com/6355099/222023068-3491e405-f195-4e11-b679-36b19a1380c7.png"> You can control the depth of the module hierarchy with the `depth` attribute (which defaults to 2). For example, if I don't limit it, this is what it outputs. <img width="366" alt="image" src="https://user-images.githubusercontent.com/6355099/222023306-3d880bb6-f534-4f98-bf10-83c4353acefc.png"> ## Other APIs FlopCounterMode(custom_mapping=...): Allows for custom flop counting functions FlopCounterMode.get_table(depth=...): Explicitly get the table as a string FlopCounterMode.flop_counts: Contains the flop information as a Dict[hierarchy: str, Dict[Op, int]] FlopCounterMode.register_hierarchy(f, name): Allows you to register additional "hierarchies" for a function. [ghstack-poisoned]

ghstack-source-id: e17f2588d5a75c6b5081505eb16d09c7ae150411 Pull Request resolved: #95751

ezyang · 2023-03-01T02:54:21Z

Who do you want to do a close code review on this?

yinghai · 2023-03-01T03:02:52Z

torch/utils/flop_counter.py

+    return flop_count
+
+
+flop_mapping = {


Can we also add support to scaled_dot_product_attention since it's more and more used nowadays?

Einsum isn’t needed since it’s a “compositeimplicit” op and decomposes into other operators. There’s an example of counting einsum flops in the tests.

I'm gonna do the attention ones in a follow-up PR.

This actually leads to an important design issue: I think there should be an API to return "unsupported" ops. Similar to https://github.com/facebookresearch/fvcore/blob/51092b5515cbb493f73de079743dd6b11cc4bbf1/fvcore/nn/jit_analysis.py#L98

The reason is that users are mainly interacting with high-level modules and ops, and aren't aware of what low-level ops are called. Imagine one day torch adds a new op for transformer, or adds a different implementation of einsum that doesn't decompose: the change should be transparent, but it makes the flop counter become a silent error!

So in fvcore we provide unsupported ops and print a warning about them by default. To make the results more meaningful, we have a list of trivial ops that are always ignored and will not appear in "unsupported ops" (https://github.com/facebookresearch/fvcore/blob/51092b5515cbb493f73de079743dd6b11cc4bbf1/fvcore/nn/jit_analysis.py#L28).

@ppwwyyxx Hmm... I think I might just add all the ops as "unsupported ops".

I think one benefit of putting it in core is that in principle, we should hopefully be able to keep the operator list up to date better.

Yeah, impl changes should be caught by tests, we have an einsum test in this PR, but not matmul?

Chillee · 2023-03-01T04:27:14Z

@ezyang Not sure I really need a close code review from anybody (in particular) 🤔 Maybe @ngimel

ppwwyyxx · 2023-03-01T07:37:49Z

torch/utils/flop_counter.py

+        func_packet = func._overloadpacket
+        if func_packet in self.flop_mapping:
+            flop_count_func = self.flop_mapping[func_packet]
+            args_shape, out_shape = tree_map(get_shape, (args, normalize_tuple(out)))


Is it a good idea to always apply get_shape on tensors? For example I wonder if there would be an op that takes a scalar boolean tensor and condition two different behaviors (with different flops) on it.

I don't think it's a big deal - if that ever shows up we can just pass through scalar tensors.

Overall, an example usage. Note that this *also* captures backwards FLOPs. ``` import torchvision.models as models import torch from torch.utils.flop_counter import FlopCounterMode inp = torch.randn(1, 3, 224, 224, device='cpu') mod = models.resnet18() flop_counter = FlopCounterMode(mod, depth=1) with flop_counter: mod(inp).sum().backward() ``` <img width="326" alt="image" src="https://user-images.githubusercontent.com/6355099/222023068-3491e405-f195-4e11-b679-36b19a1380c7.png"> You can control the depth of the module hierarchy with the `depth` attribute (which defaults to 2). For example, if I don't limit it, this is what it outputs. <img width="366" alt="image" src="https://user-images.githubusercontent.com/6355099/222023306-3d880bb6-f534-4f98-bf10-83c4353acefc.png"> ## Other APIs FlopCounterMode(custom_mapping=...): Allows for custom flop counting functions FlopCounterMode.get_table(depth=...): Explicitly get the table as a string FlopCounterMode.flop_counts: Contains the flop information as a Dict[hierarchy: str, Dict[Op, int]] FlopCounterMode.register_hierarchy(f, name): Allows you to register additional "hierarchies" for a function. [ghstack-poisoned]

ghstack-source-id: 7d04ddd18e2ff6abe41e09d6e01d03e452864982 Pull Request resolved: #95751

Chillee · 2023-03-02T18:25:35Z

@pytorchbot merge

pytorchmergebot · 2023-03-02T18:33:58Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

albanD

I need to upgrade the public API script again to pick these up...

albanD · 2023-03-02T19:56:24Z

torch/utils/flop_counter.py

+    aten.convolution_backward: conv_backward_flop,
+}
+
+def normalize_tuple(x):


Please document public API or make them private if they shouldn't be public.

ah that's what you meant

I can also just set __all__ right?

Yes but then you'll see that you can't have these functions that look public and are not in __all__. So you will have to prepend them with _ anyways.

Why will I see? It seems to work for me.

You can try python test/test_public_bindings.py, I would expect that it will fail.

Seems to work for me 🤔

Overall, an example usage. Note that this *also* captures backwards FLOPs. ``` import torchvision.models as models import torch from torch.utils.flop_counter import FlopCounterMode inp = torch.randn(1, 3, 224, 224, device='cpu') mod = models.resnet18() flop_counter = FlopCounterMode(mod, depth=1) with flop_counter: mod(inp).sum().backward() ``` <img width="326" alt="image" src="https://user-images.githubusercontent.com/6355099/222023068-3491e405-f195-4e11-b679-36b19a1380c7.png"> You can control the depth of the module hierarchy with the `depth` attribute (which defaults to 2). For example, if I don't limit it, this is what it outputs. <img width="366" alt="image" src="https://user-images.githubusercontent.com/6355099/222023306-3d880bb6-f534-4f98-bf10-83c4353acefc.png"> ## Other APIs FlopCounterMode(custom_mapping=...): Allows for custom flop counting functions FlopCounterMode.get_table(depth=...): Explicitly get the table as a string FlopCounterMode.flop_counts: Contains the flop information as a Dict[hierarchy: str, Dict[Op, int]] FlopCounterMode.register_hierarchy(f, name): Allows you to register additional "hierarchies" for a function. [ghstack-poisoned]

pytorchmergebot · 2023-03-02T20:00:45Z

The merge job was canceled. If you believe this is a mistake,then you can re trigger it through pytorch-bot.

albanD

Ok if CI is ok!

Chillee · 2023-03-02T21:06:06Z

@pytorchbot merge

pytorchmergebot · 2023-03-02T21:07:55Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

danthe3rd · 2023-03-03T14:33:14Z

torch/utils/flop_counter.py

+            def forward(ctx, *args):
+                assert(self.parents[-1] == name)
+                self.parents.pop()
+                args = tree_map(lambda x: x.clone() if isinstance(x, torch.Tensor) else x, args)


quick question: why are we calling clone() here?

If the output of an autograd.Function is a view of the input, there are some restrictions on whether you can use inplace operations or not.

Overall, an example usage. Note that this *also* captures backwards FLOPs. ``` import torchvision.models as models import torch from torch.utils.flop_counter import FlopCounterMode inp = torch.randn(1, 3, 224, 224, device='cpu') mod = models.resnet18() flop_counter = FlopCounterMode(mod, depth=1) with flop_counter: mod(inp).sum().backward() ``` <img width="326" alt="image" src="https://user-images.githubusercontent.com/6355099/222023068-3491e405-f195-4e11-b679-36b19a1380c7.png"> You can control the depth of the module hierarchy with the `depth` attribute (which defaults to 2). For example, if I don't limit it, this is what it outputs. <img width="366" alt="image" src="https://user-images.githubusercontent.com/6355099/222023306-3d880bb6-f534-4f98-bf10-83c4353acefc.png"> ## Other APIs FlopCounterMode(custom_mapping=...): Allows for custom flop counting functions FlopCounterMode.get_table(depth=...): Explicitly get the table as a string FlopCounterMode.flop_counts: Contains the flop information as a Dict[hierarchy: str, Dict[Op, int]] FlopCounterMode.register_hierarchy(f, name): Allows you to register additional "hierarchies" for a function. Pull Request resolved: pytorch/pytorch#95751 Approved by: https://github.com/ngimel, https://github.com/albanD

Overall, an example usage. Note that this *also* captures backwards FLOPs. ``` import torchvision.models as models import torch from torch.utils.flop_counter import FlopCounterMode inp = torch.randn(1, 3, 224, 224, device='cpu') mod = models.resnet18() flop_counter = FlopCounterMode(mod, depth=1) with flop_counter: mod(inp).sum().backward() ``` <img width="326" alt="image" src="https://user-images.githubusercontent.com/6355099/222023068-3491e405-f195-4e11-b679-36b19a1380c7.png"> You can control the depth of the module hierarchy with the `depth` attribute (which defaults to 2). For example, if I don't limit it, this is what it outputs. <img width="366" alt="image" src="https://user-images.githubusercontent.com/6355099/222023306-3d880bb6-f534-4f98-bf10-83c4353acefc.png"> ## Other APIs FlopCounterMode(custom_mapping=...): Allows for custom flop counting functions FlopCounterMode.get_table(depth=...): Explicitly get the table as a string FlopCounterMode.flop_counts: Contains the flop information as a Dict[hierarchy: str, Dict[Op, int]] FlopCounterMode.register_hierarchy(f, name): Allows you to register additional "hierarchies" for a function. Pull Request resolved: pytorch#95751 Approved by: https://github.com/ngimel, https://github.com/albanD

Add flop counter utility

f386ee4

[ghstack-poisoned]

Chillee added a commit that referenced this pull request Mar 1, 2023

Add flop counter utility

02f8d69

ghstack-source-id: 477bdd302c406e1c4c7e85ab9eb67dc2c9ce5da2 Pull Request resolved: #95751

github-actions bot requested review from albanD, antoniojkim, bdhirsh, ezyang, jbschlosser, miladm, SherlockNoMad, voznesenskym and wconstab March 1, 2023 01:24

Update on "Add flop counter utility"

6b7decb

[ghstack-poisoned]

Chillee added a commit that referenced this pull request Mar 1, 2023

Add flop counter utility

f3281e0

ghstack-source-id: 871593381439193148689ea823e158338cbfa154 Pull Request resolved: #95751

Update on "Add flop counter utility"

d240abc

[ghstack-poisoned]

Chillee added a commit that referenced this pull request Mar 1, 2023

Add flop counter utility

f49352d

ghstack-source-id: a43cde0a4933ebd1a1839f57c2840d7da7b94d39 Pull Request resolved: #95751

Chillee added the topic: new features topic category label Mar 1, 2023

Chillee added a commit that referenced this pull request Mar 1, 2023

Add flop counter utility

15654bc

ghstack-source-id: 59eb9a3290d67db39822b4d14d163e2b29f811cd Pull Request resolved: #95751

Chillee requested a review from ngimel March 1, 2023 01:53

Chillee added a commit that referenced this pull request Mar 1, 2023

Add flop counter utility

4b452ae

ghstack-source-id: eb8901eff886f9018a19fb5d5730d1139069064b Pull Request resolved: #95751

Chillee added the release notes: composability release notes category label Mar 1, 2023

Chillee added a commit that referenced this pull request Mar 1, 2023

Add flop counter utility

27fbd57

ghstack-source-id: e17f2588d5a75c6b5081505eb16d09c7ae150411 Pull Request resolved: #95751

Chillee mentioned this pull request Mar 1, 2023

[RFC] Flop counters in PyTorch #93121

Open

yinghai reviewed Mar 1, 2023

View reviewed changes

ppwwyyxx reviewed Mar 1, 2023

View reviewed changes

Chillee added a commit that referenced this pull request Mar 2, 2023

Add flop counter utility

cea1bfe

ghstack-source-id: 7d04ddd18e2ff6abe41e09d6e01d03e452864982 Pull Request resolved: #95751

Chillee requested review from ngimel and albanD March 2, 2023 02:48

Chillee mentioned this pull request Mar 2, 2023

Add flop formulas for sdpa #95854

Closed

ngimel approved these changes Mar 2, 2023

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Mar 2, 2023

Chillee mentioned this pull request Mar 2, 2023

Fix node provenance tracking #95901

Closed

albanD requested changes Mar 2, 2023

View reviewed changes

albanD approved these changes Mar 2, 2023

View reviewed changes

pytorchmergebot added the Merged label Mar 2, 2023

pytorchmergebot closed this in 1b1b9c8 Mar 2, 2023

carmocca mentioned this pull request Mar 3, 2023

🚀 Add FLOPs count to model summary Lightning-AI/pytorch-lightning#12567

Open

Chillee mentioned this pull request Mar 3, 2023

Provide more informative kernel names in Inductor #95940

Closed

danthe3rd reviewed Mar 3, 2023

View reviewed changes

facebook-github-bot deleted the gh/chillee/189/head branch June 8, 2023 15:54

vadimkantorov mentioned this pull request Jul 10, 2023

[proposal] "Name" string attribute for modules, parameters, buffers, tensors for more pleasant debugging (especially for graph printouts / export / studying compiled generated code) #104247

Open

This was referenced Dec 18, 2023

Could not measure the FLOPs of LSTM pytorch/tnt#624

Closed

get_module_summary assertion fail pytorch/tnt#621

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add flop counter utility #95751

Add flop counter utility #95751

Chillee commented Mar 1, 2023 •

edited

pytorch-bot bot commented Mar 1, 2023 •

edited

ezyang commented Mar 1, 2023

yinghai Mar 1, 2023

Chillee Mar 1, 2023

ppwwyyxx Mar 1, 2023

Chillee Mar 1, 2023 •

edited

Chillee Mar 2, 2023

ppwwyyxx Mar 2, 2023

Chillee Mar 2, 2023

ngimel Mar 2, 2023

Chillee commented Mar 1, 2023 •

edited

ppwwyyxx Mar 1, 2023

Chillee Mar 2, 2023

Chillee commented Mar 2, 2023

pytorchmergebot commented Mar 2, 2023

albanD left a comment

albanD Mar 2, 2023

Chillee Mar 2, 2023

Chillee Mar 2, 2023

albanD Mar 2, 2023

Chillee Mar 2, 2023

albanD Mar 2, 2023

Chillee Mar 2, 2023

pytorchmergebot commented Mar 2, 2023

albanD left a comment

Chillee commented Mar 2, 2023

pytorchmergebot commented Mar 2, 2023

danthe3rd Mar 3, 2023

Chillee Mar 3, 2023

Add flop counter utility #95751

Add flop counter utility #95751

Conversation

Chillee commented Mar 1, 2023 • edited

Other APIs

pytorch-bot bot commented Mar 1, 2023 • edited

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/95751

✅ No Failures

ezyang commented Mar 1, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Chillee Mar 1, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Chillee commented Mar 1, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Chillee commented Mar 2, 2023

pytorchmergebot commented Mar 2, 2023

Merge started

albanD left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pytorchmergebot commented Mar 2, 2023

albanD left a comment

Choose a reason for hiding this comment

Chillee commented Mar 2, 2023

pytorchmergebot commented Mar 2, 2023

Merge started

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Chillee commented Mar 1, 2023 •

edited

pytorch-bot bot commented Mar 1, 2023 •

edited

Chillee Mar 1, 2023 •

edited

Chillee commented Mar 1, 2023 •

edited