'DABNet': DABNet, 'FPENet': FPENet } if __name__ == '__main__': parser = argparse.ArgumentParser(description='ptflops sample script') parser.add_argument('--device', type=int, default=0, help='Device to store the model.') parser.add_argument('--model', choices=list(pt_models.keys()), type=str, default='ENet') parser.add_argument('--result', type=str, default=None) args = parser.parse_args() if args.result is None: ost = sys.stdout else: ost = open(args.result, 'w') with torch.cuda.device(args.device): net = pt_models[args.model](classes=19).cuda() flops, params = get_model_complexity_info(net, (3, 512, 1024), as_strings=True, print_per_layer_stat=True, ost=ost) print('Flops: ' + flops) print('Params: ' + params)
if __name__ == "__main__": device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = DF1SegX38(classes=19).to(device) summary(model, (3, 352, 480)) x = torch.randn(2, 3, 512, 1024).to(device) from fvcore.nn.jit_handles import batchnorm_flop_jit from fvcore.nn.jit_handles import generic_activation_jit supported_ops = { "aten::batch_norm": batchnorm_flop_jit, } flop_dict, _ = flop_count(model, (x, ), supported_ops) flops_count, params_count = get_model_complexity_info( model, (3, 512, 1024), as_strings=False, print_per_layer_stat=True) input = x macs, params = profile(model, inputs=(input, )) print(flop_dict) print(flops_count, params_count) print(macs, params) ''' /home/ethan/anaconda3/envs/py36_cuda101/bin/python /home/ethan/codes/Efficient-Segmentation-Networks/model/DFSegX16.py /home/ethan/anaconda3/envs/py36_cuda101/lib/python3.6/site-packages/torch/nn/functional.py:2941: UserWarning: nn.functional.upsample is deprecated. Use nn.functional.interpolate instead. warnings.warn("nn.functional.upsample is deprecated. Use nn.functional.interpolate instead.") /home/ethan/anaconda3/envs/py36_cuda101/lib/python3.6/site-packages/torch/nn/functional.py:3121: UserWarning: Default upsampling behavior when mode=bilinear is changed to align_corners=False since 0.4.0. Please specify align_corners=True if the old behavior is desired. See the documentation of nn.Upsample for details. "See the documentation of nn.Upsample for details.".format(mode)) ---------------------------------------------------------------- Layer (type) Output Shape Param # ================================================================ Conv2d-1 [-1, 3, 353, 481] 36
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # for i in [1,2,4,6,8,10,12]: for i in [16]: img_h, img_w = 32 * i, 64 * i for name, network in models.items(): print(name, network) model = network(classes=11).to(device) total_params, trainable_params, total_input_size, total_output_size, total_params_size, total_size = summary( model, (3, img_h, img_w)) # time.sleep(1) flops, params = get_model_complexity_info( model, (3, img_h, img_w), as_strings=False, print_per_layer_stat=True, ost=sys.stdout) # print('Flops: ' + flops) # print('Params: ' + params) result = dict() result['net'] = name result['size'] = '{}x{}'.format(img_h, img_w) result['Flops'] = flops / (10**9) result['params'] = params / (10**6) result['params_A'] = total_params.item() result['params_T'] = trainable_params.item() result['Input(MB)'] = total_input_size result['F/B (MB)'] = total_output_size result['Params(MB)'] = total_params_size result['Total(MB)'] = total_size