def clip_exporter_v11(op_def, context):
node, const_tensors = export_util.translate(**locals())
min_value, max_value, const_tensors = None, None, []
dtype = context.ws.FetchTensor(op_def.output[0]).dtype
for arg in op_def.arg:
if arg.name == 'low':
min_value = arg.f
elif arg.name == 'high':
max_value = arg.f
if min_value is not None:
const_tensors.append(helper.from_array(
numpy.array(min_value, dtype),
context.unique_name(op_def.input[0] + '/clip/min_value'),
))
node.input.extend([const_tensors[-1].name])
else:
node.input.extend([''])
if max_value is not None:
const_tensors.append(helper.from_array(
numpy.array(max_value, dtype),
context.unique_name(op_def.input[0] + '/clip/max_value'),
))
node.input.extend([const_tensors[-1].name])
else:
node.input.extend([''])
return node, const_tensors
def pad_exporter_v11(op_def, context):
node, pads, value = pad_exporter(**locals())
pads = helper.from_array(
numpy.array(pads, 'int64'),
context.unique_name(op_def.input[0] + '/pad/pads'),
)
value = helper.from_array(
numpy.array(value, 'float64'),
context.unique_name(op_def.input[0] + '/pad/value'),
)
node.input.extend([pads.name, value.name])
return node, [pads, value]
def slice_exporter_v10(op_def, context):
node, starts, ends = slice_exporter(**locals())
axes = helper.from_array(
numpy.arange(len(starts), dtype='int64'),
context.unique_name(op_def.input[0] + '/slice/axes'),
)
starts = helper.from_array(
numpy.array(starts, 'int64'),
context.unique_name(op_def.input[0] + '/slice/starts'),
)
ends = helper.from_array(
numpy.array(ends, 'int64'),
context.unique_name(op_def.input[0] + '/slice/ends'),
)
node.input.extend([starts.name, ends.name, axes.name])
return node, [starts, ends, axes]
def expand_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
shape = list(context.blob_shapes[op_def.output[0]])
shape = helper.from_array(
numpy.array(shape, 'int64'),
context.unique_name(op_def.input[0] + '/expand/shape'),
)
node.input.extend([shape.name])
return node, [shape]
def fill_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'Constant'
shape = list(context.blob_shapes[op_def.output[0]])
value = helper.from_array(
numpy.array(shape, 'int64'),
context.unique_name(op_def.output[0] + '/constant/value'))
helper.add_attribute(node, 'value', value)
node.ClearField('input')
return node, [value]
def top_k_exporter_v11(op_def, context):
node, (k, axis, largest, sorted) = top_k_exporter(**locals())
helper.add_attribute(node, 'axis', axis)
helper.add_attribute(node, 'largest', largest)
helper.add_attribute(node, 'sorted', sorted)
k = helper.from_array(
numpy.array([k], 'int64'),
context.unique_name(op_def.input[0] + '/top_k/k'),
)
node.input.extend([k.name])
return node, [k]
def top_k_exporter_v10(op_def, context):
node, (k, axis, largest, sorted) = top_k_exporter(**locals())
if largest == 0:
raise ValueError('TopK-10 does not support smallest mode.')
helper.add_attribute(node, 'axis', axis)
k = helper.from_array(
numpy.array([k], 'int64'),
context.unique_name(op_def.input[0] + '/top_k/k'),
)
node.input.extend([k.name])
return node, [k]
def trilu_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
k = 0
for arg in op_def.arg:
if arg.name == 'upper':
helper.add_attribute(node, 'upper', arg.i)
elif arg.name == 'k':
k = arg.i
k = helper.from_array(
numpy.array(k, 'int64'),
context.unique_name(op_def.input[0] + '/trilu/k'),
)
node.input.extend([k.name])
return node, [k]
def tile_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
repeats = []
for arg in op_def.arg:
if arg.name == 'repeats':
repeats = [e for e in arg.ints]
elif arg.name == 'repeats_desc':
repeats = helper.fetch_argument(op_def, arg, context.ws)
repeats = helper.from_array(
numpy.array(repeats, 'int64'),
context.unique_name(op_def.input[0] + '/tile/repeats'),
)
node.input.extend([repeats.name])
return node, [repeats]
def one_hot_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
helper.add_attribute(node, 'axis', -1)
depth, on_value, off_value = 1, 1., 0.
dtype = context.ws.FetchTensor(node.output[0]).dtype
for arg in op_def.arg:
if arg.name == 'depth':
depth = arg.i
elif arg.name == 'on_value':
on_value = arg.f
elif arg.name == 'off_value':
off_value = arg.f
depth = helper.from_array(
numpy.array(depth, 'int64'),
context.unique_name(op_def.input[0] + '/one_hot/depth'),
)
values = helper.from_array(
numpy.array([off_value, on_value], dtype),
context.unique_name(op_def.input[0] + '/one_hot/values'),
)
const_tensors = [depth, values]
node.input.extend([depth.name, values.name])
return node, const_tensors
def resize_v10(op_def, context):
node, const_tensors = export_util.translate(**locals())
input_shape = context.blob_shapes[op_def.input[0]]
output_shape = context.blob_shapes[op_def.output[0]]
for arg in op_def.arg:
if arg.name == 'mode':
helper.add_attribute(node, 'mode', arg.s.lower())
scales = helper.from_array(
numpy.array([
float(output_shape[i]) / input_shape[i]
for i in range(len(input_shape))
], 'float32'),
context.unique_name(op_def.input[0] + '/resize/scales'),
)
node.input.extend([scales.name])
return node, [scales]
def cumulative_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
axis = 0
for arg in op_def.arg:
if arg.name == 'axis':
axis = arg.i
elif arg.name == 'exclusive':
helper.add_attribute(node, 'exclusive', arg.i)
elif arg.name == 'reverse':
helper.add_attribute(node, 'reverse', arg.i)
axis = helper.from_array(
numpy.array(axis, 'int64'),
context.unique_name(op_def.input[0] + '/cumulative/axis'),
)
node.input.extend([axis.name])
return node, [axis]
def reshape_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
shape = dims = list(context.blob_shapes[op_def.output[0]])
for arg in op_def.arg:
if arg.name == 'dims':
dims = [int(e) for e in arg.ints]
elif arg.name == 'dims_desc':
dims = helper.fetch_argument(op_def, arg, context.ws)
for axis, dim in enumerate(dims):
shape[axis] = dim if dim <= 0 else shape[axis]
shape = helper.from_array(
numpy.array(shape, 'int64'),
context.unique_name(op_def.input[0] + '/reshape/shape'),
)
node.input.extend([shape.name])
return node, [shape]
def roi_align(op_def, context):
node, const_tensors = export_util.translate(**locals())
# Make a dummy "batch_indices".
batch_indices = helper.from_array(
numpy.zeros((context.blob_shapes[node.input[1]][0], ), 'int64'),
context.unique_name(op_def.input[0] + '/roi_align/batch_indices'),
)
node.input.extend([batch_indices.name])
for arg in op_def.arg:
if arg.name == 'pooled_h':
helper.add_attribute(node, 'output_height', arg.i)
elif arg.name == 'pooled_w':
helper.add_attribute(node, 'output_width', arg.i)
elif arg.name == 'spatial_scale':
helper.add_attribute(node, 'spatial_scale', arg.f)
elif arg.name == 'sampling_ratio':
helper.add_attribute(node, 'sampling_ratio', arg.i)
return node, [batch_indices]
def resize_v11(op_def, context):
node, const_tensors = resize_v10(**locals())
coord_mode = 'half_pixel'
for arg in op_def.arg:
if arg.name == 'mode':
if arg.s.lower() == b'nearest':
helper.add_attribute(node, 'nearest_mode', 'floor')
if arg.name == 'align_corners':
if arg.i > 0:
coord_mode = 'align_corners'
helper.add_attribute(node, 'coordinate_transformation_mode', coord_mode)
rank = len(context.blob_shapes[op_def.input[0]])
roi = helper.from_array(
numpy.array(([0] * rank + [1] * rank), 'float32'),
context.unique_name(op_def.input[0] + '/resize/roi'),
)
node.input[:] = [node.input[0], roi.name, node.input[1]]
return node, const_tensors + [roi]
def graph_def_to_onnx_graph(
cls,
graph_def,
input_names=None,
output_names=None,
input_shapes=None,
constants=None,
value_info=None,
opset_version=None,
workspace=None,
verbose=True,
):
input_names = [] if input_names is None else input_names
output_names = [] if output_names is None else output_names
constants = {} if constants is None else constants
value_info = {} if value_info is None else value_info
if not nest.is_sequence(input_names):
raise ValueError('<input_names> should be a sequence.')
if not nest.is_sequence(output_names):
raise ValueError('<output_names> should be a sequence.')
if not isinstance(constants, dict):
raise ValueError('<constants> should be a dict with name -> value.')
if not isinstance(value_info, dict):
raise ValueError('<value_info> should be a dict with name -> (dtype, shape).')
# Determine the opset version to select exporters.
if opset_version is None:
opset_version = cls._check_opset_version(opset_version)
# Create aliases for blobs.
blob_aliases = {}
for i, alias in enumerate(output_names):
blob_aliases[graph_def.output[i]] = alias
workspace.RegisterAlias(graph_def.output[i], alias)
if graph_def.output[i] in value_info:
value_info[alias] = value_info[graph_def.output[i]]
for i, alias in enumerate(input_names):
blob_aliases[graph_def.input[i]] = alias
workspace.RegisterAlias(graph_def.input[i], alias)
if graph_def.input[i] in value_info:
value_info[alias] = value_info[graph_def.input[i]]
# Maybe rewrite the input shapes for future development.
# A common case is that we should fill ``-1`` for dynamic dimension
# in the inference runtime like TensorRT.
if input_shapes is not None:
if isinstance(input_shapes, dict):
for k, v in input_shapes.items():
value_info[k] = (value_info[k][0], v)
else:
for k, v in zip(graph_def.input[:], input_shapes):
value_info[k] = (value_info[k][0], v)
# Prepare to make the graph.
onnx_graph = onnx.GraphProto(name=graph_def.name
if len(graph_def.name) > 0
else 'onnx-model')
blob_shapes, blob_names = {}, {}
blob_versions = collections.defaultdict(
int, **dict((blob_aliases.get(k, k), 1)
for k in helper.collect_inputs(graph_def)))
initializers, seen_initializers = [], set()
# Build translator context.
context = export_util.TranslatorContext(
workspace=workspace,
blob_names=blob_names,
blob_shapes=blob_shapes,
blob_versions=blob_versions,
opset_version=opset_version,
)
# Add nodes.
for op in graph_def.op:
# Get the shape of inputs and outputs.
for name in itertools.chain(op.input, op.output):
impl = workspace.GetTensor(name)
if impl is not None:
blob_shapes[name] = impl.dims
else:
blob_shapes[name] = value_info[name][1]
# Translate definition.
nodes, const_tensors = cls._make_node(op, context)
# Rewritten for names.
for node in nodes:
node.input[:] = [blob_aliases.get(e, e) for e in node.input]
node.output[:] = [blob_aliases.get(e, e) for e in node.output]
cls._rewrite_for_ssa(node, context)
# Convert constant outputs if necessary.
if None in nodes:
const_tensors = [helper.from_tensor(name, workspace)
for name in op.output]
else:
onnx_graph.node.extend(nodes)
# Merge constant tensors.
if const_tensors is not None:
value_info = {**value_info,
**dict((e.name, (e.data_type, e.dims))
for e in const_tensors)}
for tensor in const_tensors:
if tensor.name not in seen_initializers:
initializers.append(tensor)
seen_initializers.add(tensor.name)
# Add constants.
if constants is not None:
for k, v in constants.items():
initializers.append(helper.from_array(v, name=k))
# Add inputs.
for name in helper.collect_inputs(onnx_graph):
try:
onnx_graph.input.extend([
helper.make_tensor_value_info(
name=name,
elem_type=value_info[name][0],
shape=value_info[name][1])])
except KeyError:
impl = workspace.GetTensor(name)
if impl is not None:
initializer = helper.from_tensor(name, workspace)
onnx_graph.input.extend([
helper.make_tensor_value_info(
name=name,
elem_type=initializer.data_type,
shape=initializer.dims)])
if name not in seen_initializers:
initializers.append(initializer)
seen_initializers.add(initializer.name)
else:
raise ValueError(
'Info of tensor `{}` is missing, '
'specify it in <value_info>.'.format(name))
# Add initializers.
onnx_graph.initializer.extend(initializers)
# Add outputs.
onnx_graph.output.extend(
helper.make_tensor_value_info(
name=blob_names.get(name_v2, name_v2),
elem_type=value_info[name_v2][0],
shape=value_info[name_v2][1])
for name_v2 in [blob_aliases.get(name, name)
for name in set(graph_def.output)])
if verbose:
print(helper.printable_graph(onnx_graph))
return onnx_graph
