记录计算FLOP的代码。代码以VNet为例子。get_model_complexity_info可以打印每一层所需的macs和params

    # compute FLOPS & PARAMETERSfrom ptflops.flops_counter import get_model_complexity_infomodel = 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的分享