1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
|
## @package onnx
#Module caffe2.python.onnx.onnxifi
"""
ONNXIFI a Caffe2 net
"""
from caffe2.proto import caffe2_pb2
from caffe2.python import core, workspace
import caffe2.python._import_c_extension as C
import numpy as np
def onnxifi_caffe2_net(
pred_net,
input_shapes,
max_batch_size=1,
max_seq_size=1,
debug=False,
use_onnx=True,
merge_fp32_inputs_into_fp16=False,
adjust_batch=True,
block_list=None,
weight_names=None,
net_ssa_rewritten=False,
timeout=0):
"""
Transform the caffe2_net by collapsing ONNXIFI-runnable nodes into Onnxifi c2 ops
"""
shape_hints = caffe2_pb2.TensorBoundShapes()
if type(input_shapes) is caffe2_pb2.TensorBoundShapes:
shape_hints = input_shapes
elif type(input_shapes) is dict:
for k, v in input_shapes.items():
tbs = caffe2_pb2.TensorBoundShape()
tbs.name = k
tbs.shape.dims.extend(v)
tbs.dim_type.extend([caffe2_pb2.TensorBoundShape.CONSTANT] * len(tbs.shape.dims))
tbs.dim_type[0] = caffe2_pb2.TensorBoundShape.BATCH
shape_hints.shapes.extend([tbs])
shape_hints.max_batch_size = max_batch_size
shape_hints.max_feature_len = max_seq_size
pred_net_str = C.onnxifi(pred_net.SerializeToString(),
shape_hints.SerializeToString(),
block_list if block_list else [],
weight_names if weight_names is not None else [],
max_batch_size,
max_seq_size,
timeout,
adjust_batch,
debug,
merge_fp32_inputs_into_fp16,
net_ssa_rewritten,
use_onnx)
pred_net_cut = caffe2_pb2.NetDef()
pred_net_cut.ParseFromString(pred_net_str)
return pred_net_cut
|
