神经网络的flops估计
pthflops的Python项目详细描述
pytorch估计失败
简单的pytorch实用程序,用于估计给定网络的失败次数。目前只支持一些基本操作(基本上是我的模型所需的操作)。很快还会有更多。
欢迎所有捐款。
安装
您可以使用pip:
pip install pthflops
或者直接从github存储库:
git clone https://github.com/1adrianb/pytorch-estimate-flops && pytorch-estimate-flops
python setup.py install
示例
importtorchfromtorchvision.modelsimportresnet18frompthflopsimportcount_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)
忽略某些层:
importtorchfromtorchimportnnfrompthflopsimportcount_opsclassCustomLayer(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 ignoreddefforward(self,x):returnself.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'])