Repository: sovrasov/flops-counter.pytorch (language: python)

Issue:
Regarding VIT support
Hi, 

Thanks for open sourcing a useful tool! I have been using this for double-checking # of MACs and it worked well for alexnet/vgg/resnet/mobilenet/efficientnet and even convnext.

Now I would like to see whether I can use this for vit since it is supported from torchvision models (https://pytorch.org/vision/0.12/models.html), but I got some errors when I test vit_b_16 using the sample/classficiation.py -- I think in the forward hook. 

Could you please check whether it's running well from your side? If not, do you have any plan to support this as well?



Discussion / hints:
Hi. Sorry for the late response, but the current VIT implementation can't be supported by ptflops, because kwargs are used to call the forward method of MultiheadAttention module. pytorch's hooks can't get the kwargs, consequently ptflops can't acquire the shapes of the Q, K, V arguments:
```python
def forward(self, input: torch.Tensor):
        torch._assert(input.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}")
        x = self.ln_1(input)
        x, _ = self.self_attention(query=x, key=x, value=x, need_weights=False)
        x = self.dropout(x)
        x = x + input

        y = self.ln_2(x)
        y = self.mlp(y)
        return x + y
```
If we omit the kwargs, the result for vit_b_16 would be:
```
Computational complexity:       17.59 GMac
Number of parameters:           86.57 M 
```

Fix patch:
diff --git a/samples/classification.py b/samples/classification.py
index 1300505..3645ae2 100644
--- a/samples/classification.py
+++ b/samples/classification.py
@@ -2,7 +2,9 @@
 import sys
 
 import torch
-import torchvision.models as models
+import torchvision
+from packaging import version
+from torchvision import models as models
 
 from ptflops import get_model_complexity_info
 
@@ -14,6 +16,10 @@
              'densenet': models.densenet161,
              'inception': models.inception_v3}
 
+if version.parse(torchvision.__version__) > version.parse('0.14.1'):
+    pt_models['vit_b_16'] = models.vit_b_16
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='ptflops sample script')
     parser.add_argument('--device', type=int, default=0,


Test patch:
diff --git a/test_requirements.txt b/test_requirements.txt
index d4e5c50..0ba0630 100644
--- a/test_requirements.txt
+++ b/test_requirements.txt
@@ -3,3 +3,4 @@ flake8-import-order==0.18.1
 isort==4.3.21
 torchvision>=0.5.0
 pytest==7.1.2
+packaging