记录计算FLOP的代码。代码以VNet为例子。get_model_complexity_info可以打印每一层所需的macs和params
# compute FLOPS & PARAMETERS
from ptflops.flops_counter import get_model_complexity_info
model = VNet(n_channels=1, n_classes=2, normalization='batchnorm', has_dropout=False)
with torch.cuda.device(0):
macs, params = get_model_complexity_info(model, (1, 112, 112, 80), as_strings=True,
print_per_layer_stat=True, verbose=True)
print('{:<30} {:<8}'.format('Computational complexity: ', macs))
print('{:<30} {:<8}'.format('Number of parameters: ', params))
with torch.cuda.device(0):
macs, params = get_model_complexity_info(model, (1, 96, 96, 96), as_strings=True,
print_per_layer_stat=True, verbose=True)
print('{:<30} {:<8}'.format('Computational complexity: ', macs))
print('{:<30} {:<8}'.format('Number of parameters: ', params))
import ipdb; ipdb.set_trace()
有关ptflop包我放在这里:
链接:https://pan.baidu.com/s/1IcugAR1dr9etAJ1fOqEmlA?pwd=0cod
提取码:0cod
--来自百度网盘超级会员V6的分享