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,224),custom_ops={YourModule: count_your_model})

以上這篇在Pytorch中計算自己模型的FLOPs方式就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支援我們。