flops counter.pytorch

Инструмент подсчета флопов для нейронных сетей в рамках Pytorch

Этот инструмент предназначен для вычисления теоретического количества операций с умноженным добавлением в нейронных сетях. Он также может вычислить количество параметров и распечатать вычислительную стоимость данной сети.

ptflops имеет две бэкэнды, pytorch и aten . Бэкэнд pytorch - это наследие, он рассматривает только nn.Modules . Тем не менее, это все еще полезно, поскольку он обеспечивает лучшую аналитику PAR-слой для CNNS. Во всех других случаях рекомендуется использовать aten Backend, который рассматривает операции ATEN, и, следовательно, он охватывает больше модельных архитектур (включая трансформаторы). Бэкэнд по умолчанию - aten . Пожалуйста, не используйте бэкэнд pytorch для трансформаторных архитектур.

aten Backend

Рассмотрены операции:

• aten.mm, aten.matmul, aten.addmm, aten.bmm
• aten.convolution

Советы по использованию

• Используйте verbose=True чтобы увидеть операции, которые не рассматривались во время вычисления сложности.
• Этот бэкэнд печатает статистику для модулей только для модулей, непосредственно вложенных в корень nn.Module . Более глубокие модули на втором уровне гнездования не показаны в статистике для каждого слоя.
• Опция ignore_modules заставляет ptflops игнорировать перечисленные модули. Это может быть полезно для исследовательских целей. Например, можно отказаться от всех сознания из процесса подсчета, указав ignore_modules=[torch.ops.aten.convolution, torch.ops.aten._convolution] .

Бэкэнд pytorch

Поддерживаемые слои:

• Conv1d/2d/3d (включая группировку)
• Convtranspose1d/2d/3d (включая группировку)
• Batchnorm1d/2d/3d, GroupNorm, InstancEnorm1d/2d/3d, Layerorm
• Активации (Relu, Prelu, Elu, Relu6, Leakyreelu, Gelu)
• Линейный
• Пример
• Boolings (AVGPOOL1D/2D/3D, MAXPOOL1D/2D/3D и адаптивные)

Экспериментальная поддержка:

• Rnn, lstm, gru (предполагается макет NLH)
• Rnncell, lstmcell, grucell
• TORCH.NN.MultiheadAttention
• torchvision.ops.deformconv2d
• Визуальные трансформаторы из Тимма

Советы по использованию

• Этот бэкэнд не учитывает некоторые из torch.nn.functional.* И tensor.* Operations. Поэтому неподдерживаемые операции не способствуют окончательной оценке сложности. См ptflops/pytorch_ops.py:FUNCTIONAL_MAPPING,TENSOR_OPS_MAPPING для проверки поддерживаемых OPS. Иногда функциональные крючки конфликтуют с крючками для nn.Module (например, пользовательские). В этом случае подсчет этих OPS может быть отключен путем передачи backend_specific_config={"count_functional" : False} .
• ptflops запускает данную модель на случайном тензоре и оценивает количество вычислений во время вывода. Сложные модели могут иметь несколько входов, некоторые из них могут быть необязательными. Для построения нетривиального ввода можно использовать аргумент input_constructor of get_model_complexity_info . input_constructor - это функция, которая принимает входное пространственное разрешение в качестве кортежа и возвращает дикт с именованными входными аргументами модели. Далее этот дикт будет передаваться в модель как аргументы ключевого слова.
• Параметр verbose позволяет получать информацию о модулях, которые не способствуют конечным числам.
• Опция ignore_modules заставляет ptflops игнорировать перечисленные модули. Это может быть полезно для исследовательских целей. Например, можно отказаться от всех сознания из процесса подсчета, указав ignore_modules=[torch.nn.Conv2d] .

Требования

Pytorch> = 2,0. Используйте pip install ptflops==0.7.2.2 для работы с Torch 1.x.

Установите последнюю версию

От PYPI:

pip install ptflops

Из этого репозитория:

pip install --upgrade git+https://github.com/sovrasov/flops-counter.pytorch.git

Пример

 import torchvision . models as models
import torch
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 , 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 ))

Цитирование

Если PTFLOPS был полезен для вашего документа или технического отчета, пожалуйста, укажите меня:

 @online{ptflops,
  author = {Vladislav Sovrasov},
  title = {ptflops: a flops counting tool for neural networks in pytorch framework},
  year = 2018-2024,
  url = {https://github.com/sovrasov/flops-counter.pytorch},
}

Кредиты

Спасибо @warmspringwinds и Горация, он за первоначальную версию сценария.

Эталон

TOCHVISION

Модель	Входное разрешение	Params (m)	Macs (g) ( pytorch )	Mac (g) ( aten )
Алекснет	224x224	61.10	0,72	0,71
urvnext_base	224x224	88.59	15.43	15.38
Денсенет121	224x224	7,98	2.90
EfficiveNet_b0	224x224	5.29	0,41
EfficiveNet_v2_m	224x224	54.14	5.43
Googlenet	224x224	13.00	1.51
Начало_V3	224x224	27.16	5.75	5.71
maxvit_t	224x224	30,92	5.48
mnasnet1_0	224x224	4.38	0,33
mobilenet_v2	224x224	3.50	0,32
mobilenet_v3_large	224x224	5.48	0,23
regnet_y_1_6gf	224x224	11.20	1.65
Resnet18	224x224	11.69	1.83	1.81
Resnet50	224x224	25.56	4.13	4.09
resnext50_32x4d	224x224	25.03	4.29
shufflenet_v2_x1_0	224x224	2.28	0,15
Squeezenet1_0	224x224	1.25	0,84	0,82
VGG16	224x224	138.36	15.52	15.48
vit_b_16	224x224	86.57	17.61 (неправильно)	16.86
WIDE_RESNET50_2	224x224	68.88	11.45

Тимм

Модель | Входное разрешение | Params (M) | Macs (g)