1import sys as _sys 2import google.protobuf.text_format as text_format 3from six import text_type as _text_type 4 5 6 7def _convert(args): 8 if args.framework == 'caffe': 9 raise NotImplementedError("Destination [Caffe] is not implemented yet.") 10 11 elif args.framework == 'keras': 12 raise NotImplementedError("Destination [Keras] is not implemented yet.") 13 14 elif args.framework == 'tensorflow': 15 raise NotImplementedError("Destination [Tensorflow] is not implemented yet.") 16 17 elif args.framework == 'cntk': 18 raise NotImplementedError("Destination [CNTK] is not implemented yet.") 19 20 elif args.framework == 'coreml': 21 from mmdnn.conversion.coreml.coreml_emitter import CoreMLEmitter 22 assert args.inputNetwork is not None 23 assert args.inputWeight is not None 24 emitter = CoreMLEmitter(args.inputNetwork, args.inputWeight) 25 model, in_, out_ = emitter.gen_model( 26 args.inputNames, 27 args.outputNames, 28 image_input_names = set(args.imageInputNames) if args.imageInputNames else None, 29 is_bgr = args.isBGR, 30 red_bias = args.redBias, 31 blue_bias = args.blueBias, 32 green_bias = args.greenBias, 33 gray_bias = args.grayBias, 34 image_scale = args.scale, 35 class_labels = args.classInputPath if args.classInputPath else None, 36 predicted_feature_name = args.predictedFeatureName) 37 38 """ 39 from google.protobuf import text_format 40 with open(args.output+'.txt', 'w') as f: 41 f.write(text_format.MessageToString(model)) 42 """ 43 44 with open(args.output, 'wb') as f: 45 model = model.SerializeToString() 46 f.write(model) 47 48 49 return 0 50 51 elif args.framework == 'pytorch': 52 if not args.dstWeightPath or not args.IRWeightPath: 53 raise ValueError("Need to set a target weight filename.") 54 from mmdnn.conversion.pytorch.pytorch_emitter import PytorchEmitter 55 emitter = PytorchEmitter((args.IRModelPath, args.IRWeightPath)) 56 57 elif args.framework == 'mxnet': 58 from mmdnn.conversion.mxnet.mxnet_emitter import MXNetEmitter 59 if args.IRWeightPath == None: 60 emitter = MXNetEmitter(args.IRModelPath) 61 else: 62 emitter = MXNetEmitter((args.IRModelPath, args.IRWeightPath, args.inputShape, args.dstWeightPath)) 63 64 else: 65 assert False 66 67 emitter.run(args.output) 68 69 return 0 70 71 72def _get_parser(): 73 import argparse 74 75 parser = argparse.ArgumentParser(description='Convert IR model file formats to other format.') 76 77 parser.add_argument( 78 '-f', '--framework', type=_text_type, choices=['coreml'], required=True, 79 help='Format of model at srcModelPath (default is to auto-detect).' 80 ) 81 82 parser.add_argument( 83 '-in', '--inputNetwork', 84 type=_text_type, 85 required=True, 86 help='Path of the IR network architecture file.') 87 88 parser.add_argument( 89 '-iw', '--inputWeight', 90 type=_text_type, 91 required=True, 92 help='Path to the IR network weight file.') 93 94 parser.add_argument( 95 '-o', '--output', 96 type=_text_type, 97 required=True, 98 help='Path to save the destination model') 99 100 # Caffe 101 parser.add_argument( 102 '--phase', type=_text_type, choices=['train', 'test'], default='test', 103 help='[Caffe] Convert phase (train/test) for destination toolkits.' 104 ) 105 106 # For CoreML 107 parser.add_argument('--inputNames', type=_text_type, nargs='*', help='Names of the feature (input) columns, in order (required for keras models).') 108 parser.add_argument('--outputNames', type=_text_type, nargs='*', help='Names of the target (output) columns, in order (required for keras models).') 109 parser.add_argument('--imageInputNames', type=_text_type, default=[], action='append', help='Label the named input as an image. Can be specified more than once for multiple image inputs.') 110 parser.add_argument('--isBGR', action='store_true', default=False, help='True if the image data in BGR order (RGB default)') 111 parser.add_argument('--redBias', type=float, default=0.0, help='Bias value to be added to the red channel (optional, default 0.0)') 112 parser.add_argument('--blueBias', type=float, default=0.0, help='Bias value to be added to the blue channel (optional, default 0.0)') 113 parser.add_argument('--greenBias', type=float, default=0.0, help='Bias value to be added to the green channel (optional, default 0.0)') 114 parser.add_argument('--grayBias', type=float, default=0.0, help='Bias value to be added to the gray channel for Grayscale images (optional, default 0.0)') 115 parser.add_argument('--scale', type=float, default=1.0, help='Value by which the image data must be scaled (optional, default 1.0)') 116 parser.add_argument('--classInputPath', type=_text_type, default='', help='Path to class labels (ordered new line separated) for treating the neural network as a classifier') 117 parser.add_argument('--predictedFeatureName', type=_text_type, default='class_output', help='Name of the output feature that captures the class name (for classifiers models).') 118 return parser 119 120 121def _main(): 122 parser=_get_parser() 123 args = parser.parse_args() 124 ret = _convert(args) 125 _sys.exit(int(ret)) # cast to int or else the exit code is always 1 126 127 128if __name__ == '__main__': 129 _main() 130