()
| 44 | return outputs |
| 45 | |
| 46 | def main(): |
| 47 | parser = argparse.ArgumentParser( |
| 48 | description='Dump the Python Megbrain model to C++ model, by the way ' |
| 49 | 'optimizing for inference', |
| 50 | formatter_class=argparse.ArgumentDefaultsHelpFormatter |
| 51 | ) |
| 52 | parser.add_argument('input', help='input pkl model file ') |
| 53 | parser.add_argument('-o', '--output', help='output file', required=True) |
| 54 | parser.add_argument('--init-bn', action='store_true', |
| 55 | help='initialize untrained batch-normalization, to ' |
| 56 | 'avoid NaN or Inf results') |
| 57 | parser.add_argument('--silent', action='store_true', |
| 58 | help='set verbose to False in AssertEqual opr') |
| 59 | parser.add_argument('--optimize-for-inference', action='store_true', |
| 60 | help='enbale optimization for inference') |
| 61 | parser.add_argument('--discard-var-name', action='store_true', |
| 62 | help='discard variable and param names in the ' |
| 63 | 'generated output') |
| 64 | parser.add_argument('--output-strip-info', action='store_true', |
| 65 | help='output code strip information') |
| 66 | parser.add_argument('--enable-io16xc32', action='store_true', |
| 67 | help='transform the mode to float16 io float32 compute') |
| 68 | parser.add_argument('--enable-ioc16', action='store_true', |
| 69 | help='transform the dtype of the model to float16 io ' |
| 70 | 'and compute') |
| 71 | parser.add_argument('--enable-fuse-conv-bias-nonlinearity', |
| 72 | action='store_true', |
| 73 | help='fuse convolution bias and nonlinearity opr to a ' |
| 74 | 'conv_bias opr and compute') |
| 75 | parser.add_argument('--enable-hwcd4', action='store_true', |
| 76 | help='transform the model format from NCHW to NHWCD4 ' |
| 77 | 'for inference; you may need to disable CUDA and set ' |
| 78 | 'MGB_USE_MEGDNN_DBG=2') |
| 79 | parser.add_argument('--enable-nchw4', action='store_true', |
| 80 | help='transform the model format from NCHW to NCHW4 ' |
| 81 | 'for inference') |
| 82 | parser.add_argument('--enable-nchw88', action='store_true', |
| 83 | help='transform the model format from NCHW to NCHW88 ' |
| 84 | 'for inference') |
| 85 | parser.add_argument('--enable-nchw44', action='store_true', |
| 86 | help='transform the model format from NCHW to NCHW44 ' |
| 87 | 'for inference') |
| 88 | parser.add_argument('--enable-nchw44-dot', action='store_true', |
| 89 | help='transform the model format from NCHW to NCHW44_DOT ' |
| 90 | 'for optimizing armv8.2 dot in inference') |
| 91 | parser.add_argument('--enable-chwn4', action='store_true', |
| 92 | help='transform the model format to CHWN4 ' |
| 93 | 'for inference, mainly used for nvidia tensorcore') |
| 94 | parser.add_argument('--enable-nchw32', action='store_true', |
| 95 | help='transform the model format from NCHW4 to NCHW32 ' |
| 96 | 'for inference on nvidia TensoCore') |
| 97 | parser.add_argument('--enable-fuse-conv-bias-with-z', action='store_true', |
| 98 | help='fuse conv_bias with z input for inference on ' |
| 99 | 'nvidia GPU (this optimization pass will result in mismatch ' |
| 100 | 'of the precision of output of training and inference)') |
| 101 | args = parser.parse_args() |
| 102 | |
| 103 | env = FpropEnv(verbose_fprop=False) |
no test coverage detected