| 48 | |
| 49 | |
| 50 | def profile(model, inputs, verbose=True): |
| 51 | handler_collection = [] |
| 52 | |
| 53 | def add_hooks(m): |
| 54 | if len(list(m.children())) > 0: |
| 55 | return |
| 56 | |
| 57 | m.register_buffer('total_ops', torch.zeros(1)) |
| 58 | m.register_buffer('total_params', torch.zeros(1)) |
| 59 | |
| 60 | m_type = type(m) |
| 61 | fn = None |
| 62 | if m_type in register_hooks: |
| 63 | fn = register_hooks[m_type] |
| 64 | |
| 65 | if fn is None: |
| 66 | if verbose: |
| 67 | print("No implemented counting method for {} in flops_helper".format(m)) |
| 68 | else: |
| 69 | handler = m.register_forward_hook(fn) |
| 70 | handler_collection.append(handler) |
| 71 | |
| 72 | # original_device = model.parameters().__next__().device |
| 73 | training = model.training |
| 74 | |
| 75 | model.eval() |
| 76 | model.apply(add_hooks) |
| 77 | |
| 78 | # with torch.no_grad(): |
| 79 | model(*inputs) |
| 80 | |
| 81 | total_ops = 0 |
| 82 | total_params = 0 |
| 83 | for m in model.modules(): |
| 84 | if len(list(m.children())) > 0: # skip for non-leaf module |
| 85 | continue |
| 86 | total_ops += m.total_ops |
| 87 | total_params += m.total_params |
| 88 | |
| 89 | # total_ops = total_ops.item() |
| 90 | # total_params = total_params.item() |
| 91 | total_ops = total_ops[0] |
| 92 | total_params = total_params[0] |
| 93 | |
| 94 | # reset model to original status |
| 95 | model.train(training) |
| 96 | for handler in handler_collection: |
| 97 | handler.remove() |
| 98 | |
| 99 | return total_ops, total_params |
| 100 | |
| 101 | |
| 102 | multiply_adds = 1 |