在Pytorch中计算自己模型的FLOPs方式

yipeiwu_com6年前Python基础

https://github.com/Lyken17/pytorch-OpCounter

安装方法很简单:

pip install thop

基本用法:

from torchvision.models import resnet50from thop import profile
model = resnet50()
flops, params = profile(model, input_size=(1, 3, 224,224))

对自己的module进行特别的计算:

class YourModule(nn.Module):
# your definition
def count_your_model(model, x, y):
# your rule
hereflops, params = profile(model, input_size=(1, 3, 224,224),
custom_ops={YourModule: count_your_model})

以上这篇在Pytorch中计算自己模型的FLOPs方式就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持【听图阁-专注于Python设计】。

相关文章

python迭代器与生成器详解

例子 老规矩,先上一个代码: def add(s, x): return s + x def gen(): for i in range(4): yield i bas...

Python帮你微信头像任意添加装饰别再@微信官方了

Python帮你微信头像任意添加装饰别再@微信官方了

@微信官方 昨天朋友圈刷爆了@微信官方的梗,从起初的为头像添加国旗,到最后的各种Book思议的需求…而我呢?@了辣么辣么多的奥特曼,结果还是加班到12点多…最后想想,人还是得靠...

python开发之基于thread线程搜索本地文件的方法

python开发之基于thread线程搜索本地文件的方法

本文实例讲述了python开发之基于thread线程搜索本地文件的方法。分享给大家供大家参考,具体如下: 先来看看运行效果图: 利用多个线程处理搜索的问题,我们可以发现他很快.......

Python实现简单石头剪刀布游戏

Python实现简单石头剪刀布游戏

近日在学习Python的一些基础知识,觉得还是很有趣的一个一门语言!就目前的学习的一些知识,编写了一些一个简单的石头剪刀布的游戏。主要是熟悉一些Python的一些控制语句。 impo...

python实现共轭梯度法

python实现共轭梯度法

共轭梯度法是介于最速下降法与牛顿法之间的一个方法,它仅需利用一阶导数信息,但克服了最速下降法收敛慢的缺点,又避免了牛顿法需要存储和计算Hesse矩阵并求逆的缺点,共轭梯度法不仅是解决大型...