34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
'''
|
||
计算网络或者model的参数量和FLOPS（貌似是计算量）
|
||
'''
|
||
'''cr：https://blog.csdn.net/qq_35407318/article/details/109359006
|
||
http://t.csdn.cn/prmSk
|
||
|
||
'''
|
||
|
||
import torch
|
||
|
||
from thop import profile
|
||
|
||
'''计算网络的参数量和计算量'''
|
||
|
||
if __name__ == "__main__":
|
||
# 需要使用device来指定网络在GPU还是CPU运行
|
||
from taihuyuan_pv.compared_experiment.deeplabv3Plus.model.modeling import deeplabv3plus_resnet50
|
||
|
||
|
||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
|
||
model= deeplabv3plus_resnet50(num_classes=2,output_stride=16,pretrained_backbone=False).to(device)
|
||
|
||
# model = baselinemodel_resnet50(modelname="baseline32_ms", num_classes=2, output_stride=16, pretrained_backbone=False).to(device)
|
||
|
||
input = torch.randn(1, 3, 512, 512).to(device)
|
||
flops, params = profile(model, inputs=(input,))
|
||
print("FLOPS:",flops)
|
||
print("params:",params)
|
||
|
||
print("Total FLOPS: %.2fGflops" % (flops/1e9))
|
||
print("Total params: %.2fM" % (params/1e6))
|
||
|
||
# torch.save(model.state_dict(),"test1.pth") |