| |
| |
|
|
|
|
|
|
|
|
|
|
| import json |
|
|
| from caffe2.proto import caffe2_pb2 |
| import click |
| from onnx import ModelProto |
|
|
| from caffe2.python.onnx.backend import Caffe2Backend as c2 |
| import caffe2.python.onnx.frontend as c2_onnx |
|
|
|
|
| @click.command( |
| help='convert caffe2 net to onnx model', |
| context_settings={ |
| 'help_option_names': ['-h', '--help'] |
| } |
| ) |
| @click.argument('caffe2_net', type=click.File('rb')) |
| @click.option('--caffe2-net-name', |
| type=str, |
| help="Name of the caffe2 net") |
| @click.option('--caffe2-init-net', |
| type=click.File('rb'), |
| help="Path of the caffe2 init net pb file") |
| @click.option('--value-info', |
| type=str, |
| help='A json string providing the ' |
| 'type and shape information of the inputs') |
| @click.option('-o', '--output', required=True, |
| type=click.File('wb'), |
| help='Output path for the onnx model pb file') |
| def caffe2_to_onnx(caffe2_net, |
| caffe2_net_name, |
| caffe2_init_net, |
| value_info, |
| output): |
| c2_net_proto = caffe2_pb2.NetDef() |
| c2_net_proto.ParseFromString(caffe2_net.read()) |
| if not c2_net_proto.name and not caffe2_net_name: |
| raise click.BadParameter( |
| 'The input caffe2 net does not have name, ' |
| '--caffe2-net-name must be provided') |
| c2_net_proto.name = caffe2_net_name or c2_net_proto.name |
| if caffe2_init_net: |
| c2_init_net_proto = caffe2_pb2.NetDef() |
| c2_init_net_proto.ParseFromString(caffe2_init_net.read()) |
| c2_init_net_proto.name = '{}_init'.format(caffe2_net_name) |
| else: |
| c2_init_net_proto = None |
|
|
| if value_info: |
| value_info = json.loads(value_info) |
|
|
| onnx_model = c2_onnx.caffe2_net_to_onnx_model( |
| predict_net=c2_net_proto, |
| init_net=c2_init_net_proto, |
| value_info=value_info) |
|
|
| output.write(onnx_model.SerializeToString()) |
|
|
|
|
| @click.command( |
| help='convert onnx model to caffe2 net', |
| context_settings={ |
| 'help_option_names': ['-h', '--help'] |
| } |
| ) |
| @click.argument('onnx_model', type=click.File('rb')) |
| @click.option('-o', '--output', required=True, |
| type=click.File('wb'), |
| help='Output path for the caffe2 net file') |
| @click.option('--init-net-output', |
| required=True, |
| type=click.File('wb'), |
| help='Output path for the caffe2 init net file') |
| def onnx_to_caffe2(onnx_model, output, init_net_output): |
| onnx_model_proto = ModelProto() |
| onnx_model_proto.ParseFromString(onnx_model.read()) |
|
|
| init_net, predict_net = c2.onnx_graph_to_caffe2_net(onnx_model_proto) |
| init_net_output.write(init_net.SerializeToString()) |
| output.write(predict_net.SerializeToString()) |
|
|
