def caffe2_op_to_onnx_node(cls, op_def, shapes):
if C.support_onnx_export(op_def.type):
shape_list = list(shapes.values())
node_strs, tensor_strs = C.export_to_onnx(
op_def.SerializeToString(), shapes)
nodes = []
for s in node_strs:
node = NodeProto()
node.ParseFromString(s)
nodes.append(node)
const_tensors = []
for s in tensor_strs:
tensor = TensorProto()
tensor.ParseFromString(s)
const_tensors.append(tensor)
return nodes, const_tensors
elif op_def.type in cls._special_operators:
translator = getattr(cls, cls._special_operators[op_def.type])
else:
translator = cls._common_caffe2_op_to_onnx_node
nodes = translator(op_def, shapes)
const_tensors = []
if isinstance(nodes, tuple):
nodes, const_tensors = nodes
if not isinstance(nodes, collections.Iterable):
nodes = [nodes]
return nodes, const_tensors
def check_node() -> None:
parser = argparse.ArgumentParser('check-node')
parser.add_argument('node_pb', type=argparse.FileType('rb'))
args = parser.parse_args()
node = NodeProto()
node.ParseFromString(args.node_pb.read())
checker.check_node(node)
