| 17 | logger = get_logger(__name__) |
| 18 | |
| 19 | def optimize_for_inference(args, outputs): |
| 20 | args_map = { |
| 21 | 'enable_io16xc32': 'f16_io_f32_comp', |
| 22 | 'enable_ioc16': 'f16_io_comp', |
| 23 | 'enable_hwcd4': 'use_nhwcd4', |
| 24 | 'enable_nchw4': 'use_nchw4', |
| 25 | 'enable_nchw88': 'use_nchw88', |
| 26 | 'enable_nchw44': 'use_nchw44', |
| 27 | 'enable_nchw44_dot': 'use_nchw44_dot', |
| 28 | 'enable_nchw32': 'use_nchw32', |
| 29 | 'enable_chwn4': 'use_chwn4', |
| 30 | 'enable_fuse_conv_bias_nonlinearity': 'fuse_conv_bias_nonlinearity', |
| 31 | 'enable_fuse_conv_bias_with_z': 'fuse_conv_bias_with_z', |
| 32 | } |
| 33 | kwargs = {} |
| 34 | for k, v in args_map.items(): |
| 35 | if getattr(args, k): |
| 36 | assert args.optimize_for_inference, ( |
| 37 | 'optimize_for_inference should be set when {} is given'.format( |
| 38 | k)) |
| 39 | kwargs[v] = True |
| 40 | |
| 41 | if args.optimize_for_inference: |
| 42 | return mgb.optimize_for_inference(outputs, **kwargs) |
| 43 | |
| 44 | return outputs |
| 45 | |
| 46 | def main(): |
| 47 | parser = argparse.ArgumentParser( |