简单的pytorch实用程序，用于估算给定网络的FLOP数量

简单的pytorch实用程序，用于估算给定网络的FLOP数量

pytorch-estimate-flops

Simple pytorch utility that estimates the number of FLOPs for a given network. For now only some basic operations are supported (basically the ones I needed for my models). More will be added soon.

All contributions are welcomed.

Installation

You can install the model using pip:

pip install pthflops

or directly from the github repository:

git clone https://github.com/1adrianb/pytorch-estimate-flops && pytorch-estimate-flopspython setup.py install

Example

import torchfrom torchvision.models import resnet18from pthflops import count_ops# Create a network and a corresponding inputdevice = 'cuda:0'model = resnet18().to(device)inp = torch.rand(1,3,224,224).to(device)# Count the number of FLOPscount_ops(model, inp)

Ignoring certain layers:

import torchfrom torch import nnfrom pthflops import count_opsclass CustomLayer(nn.Module): def __init__(self): super(CustomLayer, self).__init__() self.conv1 = nn.Conv2d(5, 5, 1, 1, 0) # ... other layers present inside will also be ignored def forward(self, x): return self.conv1(x)# Create a network and a corresponding inputinp = torch.rand(1,5,7,7)net = nn.Sequential( nn.Conv2d(5, 5, 1, 1, 0), nn.ReLU(inplace=True), CustomLayer())# Count the number of FLOPscount_ops(net, inp, ignore_layers=['CustomLayer'])

版权声明：本文内容由网络用户投稿，版权归原作者所有，本站不拥有其著作权，亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容，请联系我们jiasou666@gmail.com 处理，核实后本网站将在24小时内删除侵权内容。