Update readme

README.md

@@ -2,11 +2,26 @@
 [![Pypi version](https://img.shields.io/pypi/v/ptflops.svg)](https://pypi.org/project/ptflops/)
 [![Build Status](https://travis-ci.com/sovrasov/flops-counter.pytorch.svg?branch=master)](https://travis-ci.com/sovrasov/flops-counter.pytorch)
 
-This script is designed to compute the theoretical amount of multiply-add operations
-in convolutional neural networks. It can also compute the number of parameters and
+This tool is designed to compute the theoretical amount of multiply-add operations
+in neural networks. It can also compute the number of parameters and
 print per-layer computational cost of a given network.
 
-Supported layers:
+`ptflops` has two backends, `pytorch` and `aten`. `pytorch` backend is a legacy one, it considers `nn.Modules` only. However,
+it's still useful, since it provides a better par-layer analytics for CNNs. In all other cases it's recommended to use
+`aten` backend, which considers aten operations, and therefore it covers more model architectures (including transformers).
+
+## `aten` backend
+### Operations considered:
+- aten.mm, aten.matmul, aten.addmm, aten.bmm
+- aten.convolution
+
+### Usage tips
+- Use `verbose=True` to see the operations which were not considered during complexity computation.
+- This backend prints per-module statistics only for modules directly nested into the root `nn.Module`.
+Deeper modules at the second level of nesting are not shown in the per-layer statistics.
+
+## `pytorch` backend
+### Supported layers:
 - Conv1d/2d/3d (including grouping)
 - ConvTranspose1d/2d/3d (including grouping)
 - BatchNorm1d/2d/3d, GroupNorm, InstanceNorm1d/2d/3d, LayerNorm
@@ -22,20 +37,20 @@ Experimental support:
 - torchvision.ops.DeformConv2d
 - visual transformers from [timm](/~https://github.com/huggingface/pytorch-image-models)
 
-Requirements: Pytorch >= 1.1, torchvision >= 0.3
+### Usage tips
 
-Thanks to @warmspringwinds for the initial version of script.
-
-## Usage tips
-
-- This tool doesn't take into account some of the `torch.nn.functional.*` and `tensor.*` operations. Therefore unsupported operations are
+- This backend doesn't take into account some of the `torch.nn.functional.*` and `tensor.*` operations. Therefore unsupported operations are
 not contributing to the final complexity estimation. See `ptflops/pytorch_ops.py:FUNCTIONAL_MAPPING,TENSOR_OPS_MAPPING` to check supported ops.
 - `ptflops` launches a given model on a random tensor and estimates amount of computations during inference. Complicated models can have several inputs, some of them could be optional. To construct non-trivial input one can use the `input_constructor` argument of the `get_model_complexity_info`. `input_constructor` is a function that takes the input spatial resolution as a tuple and returns a dict with named input arguments of the model. Next this dict would be passed to the model as a keyword arguments.
 - `verbose` parameter allows to get information about modules that don't contribute to the final numbers.
 - `ignore_modules` option forces `ptflops` to ignore the listed modules. This can be useful
 for research purposes. For instance, one can drop all convolutions from the counting process
 specifying `ignore_modules=[torch.nn.Conv2d]`.
 
+Requirements: Pytorch >= 1.1, torchvision >= 0.3
+
+Thanks to @warmspringwinds and Horace He for the initial version of the script.
+
 ## Install the latest version
 From PyPI:
 ```bash
@@ -55,7 +70,12 @@ from ptflops import get_model_complexity_info
 
 with torch.cuda.device(0):
   net = models.densenet161()
-  macs, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True,
+  macs, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True, backend='pytorch'
+                                           print_per_layer_stat=True, verbose=True)
+  print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
+  print('{:<30}  {:<8}'.format('Number of parameters: ', params))
+
+  macs, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True, backend='aten'
                                            print_per_layer_stat=True, verbose=True)
   print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
   print('{:<30}  {:<8}'.format('Number of parameters: ', params))
@@ -67,7 +87,7 @@ If ptflops was useful for your paper or tech report, please cite me:
 @online{ptflops,
   author = {Vladislav Sovrasov},
   title = {ptflops: a flops counting tool for neural networks in pytorch framework},
-  year = 2018-2023,
+  year = 2018-2024,
   url = {/~https://github.com/sovrasov/flops-counter.pytorch},
 }
 ```