| 73 | return sum([reduce(operator.mul, i.size(), 1) for i in model.parameters()]) |
| 74 | |
| 75 | def measure_model(model, input_image_size, forward_param=None): |
| 76 | flop_counts = [] |
| 77 | param_counts = [] |
| 78 | multi_add = 2 |
| 79 | |
| 80 | # loop over all model parts |
| 81 | for m in model.modules(): |
| 82 | if isinstance(m, nn.Conv2d): |
| 83 | def hook(module, x): |
| 84 | out_h = int((x[0].size()[2] + 2 * module.padding[0] - module.kernel_size[0]) / module.stride[0] + 1) |
| 85 | out_w = int((x[0].size()[3] + 2 * module.padding[1] - module.kernel_size[1]) / module.stride[1] + 1) |
| 86 | ops = module.in_channels * module.out_channels * module.kernel_size[0] * module.kernel_size[1] * out_h * out_w / module.groups * multi_add |
| 87 | flop_counts.append(ops) |
| 88 | param_counts.append(get_layer_param(module)) |
| 89 | m.register_forward_pre_hook(hook) |
| 90 | |
| 91 | elif isinstance(m, nn.ReLU) or isinstance(m, nn.PReLU): |
| 92 | def hook(module, x): |
| 93 | ops = x[0].numel() |
| 94 | flop_counts.append(ops) |
| 95 | param_counts.append(get_layer_param(module)) |
| 96 | m.register_forward_pre_hook(hook) |
| 97 | |
| 98 | elif isinstance(m, nn.AvgPool2d): |
| 99 | def hook(module, x): |
| 100 | in_w = x[0].size()[2] |
| 101 | kernel_ops = module.kernel_size * module.kernel_size |
| 102 | out_w = int((in_w + 2 * module.padding - module.kernel_size) / module.stride + 1) |
| 103 | out_h = int((in_w + 2 * module.padding - module.kernel_size) / module.stride + 1) |
| 104 | ops = x[0].size()[0] * x[0].size()[1] * out_w * out_h * kernel_ops |
| 105 | param_counts.append(get_layer_param(module)) |
| 106 | m.register_forward_pre_hook(hook) |
| 107 | |
| 108 | elif isinstance(m, nn.AdaptiveAvgPool2d): |
| 109 | def hook(module, x): |
| 110 | ops = x[0].size()[0] * x[0].size()[1] * x[0].size()[2] * x[0].size()[3] |
| 111 | flop_counts.append(ops) |
| 112 | param_counts.append(get_layer_param(module)) |
| 113 | m.register_forward_pre_hook(hook) |
| 114 | |
| 115 | elif isinstance(m, nn.Linear): |
| 116 | def hook(module, x): |
| 117 | weight_ops = module.weight.numel() * multi_add |
| 118 | bias_ops = module.bias.numel() |
| 119 | ops = x[0].size()[0] * (weight_ops + bias_ops) |
| 120 | flop_counts.append(ops) |
| 121 | param_counts.append(get_layer_param(module)) |
| 122 | m.register_forward_pre_hook(hook) |
| 123 | |
| 124 | elif isinstance(m, nn.BatchNorm2d) or isinstance(m, SyncBatchNorm2d) or isinstance(m, nn.BatchNorm1d) \ |
| 125 | or isinstance(m, nn.Dropout2d) or isinstance(m, nn.Dropout): |
| 126 | def hook(module, x): |
| 127 | param_counts.append(get_layer_param(module)) |
| 128 | m.register_forward_pre_hook(hook) |
| 129 | |
| 130 | else: |
| 131 | # print('unknown layer type: %s' % type(m)) |
| 132 | pass |