| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788 |
- ## @package onnx
- # Module caffe2.python.onnx.bin.conversion
- 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())
|
