If you want to thest former's flop, you need to give default value to inputs in model.forward(), the following code can only pass one argument to forward()
(model,x_shape)
| 102 | plt.savefig(name, bbox_inches='tight') |
| 103 | |
| 104 | def test_params_flop(model,x_shape): |
| 105 | """ |
| 106 | If you want to thest former's flop, you need to give default value to inputs in model.forward(), the following code can only pass one argument to forward() |
| 107 | """ |
| 108 | model_params = 0 |
| 109 | for parameter in model.parameters(): |
| 110 | model_params += parameter.numel() |
| 111 | print('INFO: Trainable parameter count: {:.2f}M'.format(model_params / 1000000.0)) |
| 112 | from ptflops import get_model_complexity_info |
| 113 | with torch.cuda.device(0): |
| 114 | macs, params = get_model_complexity_info(model.cuda(), x_shape, as_strings=True, print_per_layer_stat=True) |
| 115 | # print('Flops:' + flops) |
| 116 | # print('Params:' + params) |
| 117 | print('{:<30} {:<8}'.format('Computational complexity: ', macs)) |
| 118 | print('{:<30} {:<8}'.format('Number of parameters: ', params)) |