Pytorch中计算自己模型的FLOPs(profile(),clever_format())

39 篇文章 9 订阅
18 篇文章 1 订阅

提到计算量必然涉及FLOPS, FLOPs和MACs,并不是同样的意思,切忌混淆。

FLOPS:Floating Point Operations Per Second,每秒浮点运算次数,是一个衡量硬件计算性能的指标。

FLOPs:Floating Point Operations,浮点运算次数,常用来衡量模型的计算复杂度,即计算量。(大多数论文会在表格中给出模型的FLOPs指标。不过,该指标与实际速度并不完全正相关,我自己测也是这样,所以也有一些论文会给FPS或者1/FPS即单图片推理时间。)

显然,FLOPS和FLOPs不是一回事,FLOPS的积分可以理解为等于FLOPs。

MACs:Multiply–Accumulate Operations,1 MAC = 1次加法+1次乘法操作。

因为,每次运算就等于1次FLOPs,所以,近似地,从数值上看我们有

在这里插入图片描述

安装thop

pip install thop

基础用法

以查看resnet50的FLOPs为例

from torchvision.models import resnet50
from thop import profile
model = resnet50()
input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input, ))

查看自己模型的FLOPs

class YourModule(nn.Module):
    # your definition
def count_your_model(model, x, y):
    # your rule here
 
input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input, ), 
                        custom_ops={YourModule: count_your_model})

提升输出结果的可读性,clever_format是更好的输出形式。
调用thop.clever_format

from thop import clever_format
flops, params = clever_format([flops, params], "%.3f")
  • 0
    点赞
  • 6
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
可以使用下面的代码计算PyTorch模型FLOPs(浮点操作次数): ```python import torch from torch.autograd import Variable def print_model_parm_flops(model, input_size, custom_layers): multiply_adds = 1 params = 0 flops = 0 input = Variable(torch.rand(1, *input_size)) def register_hook(module): def hook(module, input, output): class_name = str(module.__class__).split(".")[-1].split("'")[0] if class_name == 'Conv2d': out_h, out_w = output.size()[2:] kernel_h, kernel_w = module.kernel_size in_channels = module.in_channels out_channels = module.out_channels if isinstance(module.padding, int): pad_h = pad_w = module.padding else: pad_h, pad_w = module.padding if isinstance(module.stride, int): stride_h = stride_w = module.stride else: stride_h, stride_w = module.stride params += out_channels * (in_channels // module.groups) * kernel_h * kernel_w flops += out_channels * (in_channels // module.groups) * kernel_h * kernel_w * out_h * out_w / (stride_h * stride_w) elif class_name == 'Linear': weight_flops = module.weight.nelement() * input[0].nelement() // module.weight.size(1) bias_flops = module.bias.nelement() flops = weight_flops + bias_flops params = weight_flops + bias_flops elif class_name in custom_layers: custom_flops, custom_params = custom_layers[class_name](module, input, output) flops += custom_flops params += custom_params else: print(f"Warning: module {class_name} not implemented") if not isinstance(module, torch.nn.Sequential) and \ not isinstance(module, torch.nn.ModuleList) and \ not (module == model): hooks.append(module.register_forward_hook(hook)) # loop through the model architecture and register hooks for each layer hooks = [] model.apply(register_hook) # run the input through the model model(input) # remove the hooks for hook in hooks: hook.remove() print(f"Number of parameters: {params}") print(f"Number of FLOPs: {flops}") return flops, params ``` 调用这个函数需要传入模型、输入大小和一个自定义图层字典,其字典的键是自定义层的名称,值是一个函数,该函数接受模块,输入和输出,返回FLOPs和参数数量。例如,如果您的模型包含一个名为MyLayer的自定义层,则可以将以下内容添加到字典: ```python def my_layer_impl(module, input, output): # compute FLOPs and params for MyLayer flops = ... params = ... return flops, params custom_layers = {'MyLayer': my_layer_impl} ``` 使用示例: ```python import torchvision.models as models model = models.resnet18() input_size = (3, 224, 224) custom_layers = {} flops, params = print_model_parm_flops(model, input_size, custom_layers) ``` 该函数将打印出模型的参数数量和FLOPs

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值