()
| 18 | |
| 19 | |
| 20 | def parse_args(): |
| 21 | parser = argparse.ArgumentParser(description='count model flops') |
| 22 | parser.add_argument('config', help='config file path') |
| 23 | parser.add_argument( |
| 24 | '--repeat_num', default=10, type=int, help='repeat number') |
| 25 | args = parser.parse_args() |
| 26 | return args |
| 27 | |
| 28 | |
| 29 | def flatten_inputs(model, inputs): |