深度学习计算模型FLOPS的代码

记录计算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的分享