def _convert_softmax_impl(op_type, func, input_names, output_names):
axis = func.axis
ndim = len(func.inputs[0].shape)
if axis == ndim - 1:
return onnx_helper.make_node(op_type,
input_names,
output_names,
axis=axis),
# Chainer's softmax computes the softmax along a single axis while
# ONNX's computes along the specified axis and all axes after the
# specified axis. To emulate Chainer's by ONNX's, we transpose the
# single specified axis to the last axis, compute the softmax, and
# transpose back to the original shape.
gb = onnx_helper.GraphBuilder()
perm = list(range(ndim))
perm[axis], perm[-1] = perm[-1], perm[axis]
transposed = gb.op('Transpose', input_names, perm=perm)
softmaxed = gb.op(op_type, [transposed], axis=ndim - 1)
gb.op('Transpose', [softmaxed], perm=perm)
return gb.nodes(output_names=output_names)
def convert_Repeat(func, opset_version, input_names, output_names, context):
repeats = func.repeats
if len(repeats) > 1:
raise NotImplementedError(
'ONNX-Chainer currently does not support elementwise repeat')
gb = onnx_helper.GraphBuilder()
inputs = list(input_names)
axis = func.axis
if axis is None:
shape_name = context.add_const(np.array([-1], dtype=np.int64), 'shape')
input_names.append(shape_name)
inputs = [gb.op('Reshape', input_names)]
scales = [float(repeats[0])]
else:
scales = [1.0] * func.inputs[0].data.ndim
scales[axis] = float(repeats[0])
if opset_version == 7:
gb.op_output_named('Upsample', inputs, output_names, scales=scales)
return gb.nodes()
scales_name = context.add_const(np.array(scales, dtype=np.float32),
'scales')
if opset_version in [9, 10]:
inputs.append(scales_name)
op = 'Upsample' if opset_version == 9 else 'Resize'
gb.op_output_named(op, inputs, output_names)
return gb.nodes()
if opset_version == 11:
roi = context.add_const(np.array([]), 'roi')
inputs.extend([roi, scales_name])
gb.op_output_named('Resize', inputs, output_names)
return gb.nodes()
def convert_SoftmaxCrossEntropy(
func, opset_version, input_names,
output_names, context, parameters):
# obtain input variable
if not isinstance(func, chainer.FunctionNode):
raise NotImplementedError(
'SoftmaxCrossEntropy is currently supported for Chainer>=6.0.0a1.')
x_var, t_var = func.get_retained_inputs()
if len(x_var.shape) != 2:
raise NotImplementedError(
'ONNX-Chainer currently handles SoftmaxCrossEntropy only when '
'the dimension of input variable x is exactly two.')
if np.any(t_var.array == func.ignore_label):
raise NotImplementedError(
'ONNX-Chainer currently handles SoftmaxCrossEntropy only when '
'ignore_label is not used in input variable t.')
if (not func.normalize) or (func.class_weight is not None) or\
(func.ignore_label != -1) or (func.reduce != 'mean'):
raise NotImplementedError(
'ONNX-Chainer currently handles SoftmaxCrossEntropy only when '
'argument parameters are default setting.')
# create intermediate values
gb = onnx_helper.GraphBuilder()
x, t = input_names
y_log = gb.op('LogSoftmax', [x])
depth = gb.const(np.array([x_var.shape[1]], dtype=np.int32))
zeroone = gb.const(np.array([0, 1], dtype=x_var.dtype))
th = gb.op('OneHot', [t, depth, zeroone])
s0 = gb.op('Mul', [y_log, th])
sn = gb.op('Neg', [s0])
sr = gb.op('ReduceSum', [sn], axes=[1], keepdims=0)
gb.op_output_named('ReduceMean', [sr], output_names, axes=[0], keepdims=0)
return gb.nodes()
def convert_GetItem(func, opset_version, input_names, output_names, context):
x = func.inputs[0]
axes, starts, ends = [], [], []
squeeze_idxs, unsqueeze_idxs = [], []
skipped = 0 # when set ellipsis, need to skip index rolling
prev_gathered_axis = -1
gather_axis = -1
gather_idx = None # when GatherND, set first array for broadcasting
gather_nd_idx = None
is_used_slice_whole = False # GatherND does not support axis, need to care
for i, idx in enumerate(func.slices):
# axis means the index of input x, adjust None and Ellipsis counts
axis = i - len(unsqueeze_idxs) + skipped
if isinstance(idx, slice):
if idx.step is not None and idx.step != 1:
raise ValueError(
'GetItem with {}step slicing is not supported in ONNX '
'Slice operator'.format(idx.step))
if idx.start is None and idx.stop is None:
is_used_slice_whole = True
continue
axes.append(axis)
starts.append(0 if idx.start is None else idx.start)
ends.append(x.shape[axis] if idx.stop is None else idx.stop)
elif isinstance(idx, int):
axes.append(axis)
starts.append(idx)
ends.append(idx + 1)
squeeze_idxs.append(axis)
elif isinstance(idx, np.ndarray) and idx.ndim == 0:
scalar_idx = idx.item()
axes.append(axis)
starts.append(scalar_idx)
ends.append(scalar_idx + 1)
squeeze_idxs.append(axis)
elif idx is None:
unsqueeze_idxs.append(i - len(squeeze_idxs) + skipped)
elif idx is Ellipsis:
# calculate rest slice number except None, GetItem does not allow
# multiple Ellipsis, so ignore latter Ellipsis count
rest_slice_len = len(
[idx_ for idx_ in func.slices[i + 1:] if idx_ is not None])
assert skipped == 0
skipped = len(x.shape) - axis - rest_slice_len - 1
elif isinstance(idx, (list, ) + chainer.get_array_types()):
if prev_gathered_axis >= 0:
if (i - 1) != prev_gathered_axis:
raise ValueError(
'ONNX-Chainer does not support non-consecutive'
'multiple advanced indexing')
if is_used_slice_whole:
raise ValueError(
'ONNX-Chainer does not support whole indexing(`[:]`)'
'in front of multiple advanced indexing')
if unsqueeze_idxs:
raise ValueError(
'ONNX-Chainer does not support new axis in front of '
'multiple advanced indexing')
# multiple advanced index, convert to GatherND
idx_array = _to_ndarray(idx)
base_idx = gather_idx if gather_nd_idx is None else\
gather_nd_idx
gather_nd_idx = np.vstack((base_idx, idx_array))
prev_gathered_axis = i
else:
# convert to Gather, if next index is also list, change to
# GatherND
gather_axis = axis - len(squeeze_idxs) + len(unsqueeze_idxs)
gather_idx = _to_ndarray(idx)
prev_gathered_axis = i
else:
raise ValueError(
'GetItem with type {} cannot handle in ONNX Slice, so that '
'ONNX-Chainer does not accept the type'.format(type(idx)))
gb = onnx_helper.GraphBuilder()
slice_output = input_names
if axes:
output = get_slice_node(gb, opset_version, context, slice_output, axes,
starts, ends)
slice_output = [output]
if squeeze_idxs:
output = gb.op('Squeeze', slice_output, axes=squeeze_idxs)
slice_output = [output]
if unsqueeze_idxs:
output = gb.op('Unsqueeze', slice_output, axes=unsqueeze_idxs)
slice_output = [output]
if gather_nd_idx is not None:
gather_nd_idx_name = context.add_const(gather_nd_idx.T, 'indices')
slice_output.append(gather_nd_idx_name)
gb.op('GatherND', slice_output)
elif gather_idx is not None:
gather_idx_name = context.add_const(gather_idx, 'indices')
slice_output.append(gather_idx_name)
gb.op('Gather', slice_output, axis=gather_axis)
return gb.nodes(output_names=output_names)
def convert_GetItem(func, opset_version, input_names, output_names, context):
x = func.inputs[0]
axes, starts, ends = [], [], []
squeeze_idxs, unsqueeze_idxs = [], []
skipped = 0 # when set ellipsis, need to skip index rolling
gather_axis, gather_idx = [], []
for i, idx in enumerate(func.slices):
# axis means the index of input x, adjust None and Ellipsis counts
axis = i - len(unsqueeze_idxs) + skipped
if isinstance(idx, slice):
if idx.step is not None and idx.step != 1:
raise ValueError(
'GetItem with {}step slicing is not supported in ONNX '
'Slice operator'.format(idx.step))
if idx.start is None and idx.stop is None:
continue
axes.append(axis)
starts.append(0 if idx.start is None else idx.start)
ends.append(x.shape[axis] if idx.stop is None else idx.stop)
elif isinstance(idx, int):
axes.append(axis)
starts.append(idx)
ends.append(idx + 1)
squeeze_idxs.append(axis)
elif isinstance(idx, np.ndarray) and idx.ndim == 0:
scalar_idx = idx.item()
axes.append(axis)
starts.append(scalar_idx)
ends.append(scalar_idx + 1)
squeeze_idxs.append(axis)
elif idx is None:
unsqueeze_idxs.append(i - len(squeeze_idxs) + skipped)
elif idx is Ellipsis:
# calculate rest slice number except None, GetItem does not allow
# multiple Ellipsis, so ignore latter Ellipsis count
rest_slice_len = len(
[idx_ for idx_ in func.slices[i + 1:] if idx_ is not None])
assert skipped == 0
skipped = len(x.shape) - axis - rest_slice_len - 1
elif isinstance(idx, (list, ) + chainer.get_array_types()):
if gather_axis:
raise ValueError(
'ONNX-Chainer does not support multiple advanced index')
gather_axis.append(axis - len(squeeze_idxs) + len(unsqueeze_idxs))
if isinstance(idx, list):
gather_idx = np.array(idx, dtype=np.int64)
else:
gather_idx = chainer.cuda.to_cpu(idx)
else:
# not support advanced index like `array[[0,1], [0, 1]]`
raise ValueError(
'GetItem with type {} cannot handle in ONNX Slice, so that '
'ONNX-Chainer does not accept the type'.format(type(idx)))
gb = onnx_helper.GraphBuilder()
slice_output = input_names
if axes:
output = get_slice_node(gb, opset_version, context, slice_output, axes,
starts, ends)
slice_output = [output]
if squeeze_idxs:
output = gb.op('Squeeze', slice_output, axes=squeeze_idxs)
slice_output = [output]
if unsqueeze_idxs:
output = gb.op('Unsqueeze', slice_output, axes=unsqueeze_idxs)
slice_output = [output]
if gather_axis:
gather_idx_name = context.add_const(gather_idx, 'indices')
slice_output.append(gather_idx_name)
gb.op('Gather', slice_output, axis=gather_axis[0])
return gb.nodes(output_names=output_names)
def convert_RsqrtGPU(func, opset_version, input_names, output_names, context):
gb = onnx_helper.GraphBuilder()
sqrt_out = gb.op('Sqrt', input_names)
gb.op('Reciprocal', [sqrt_out])
return gb.nodes(output_names)
def convert_n_step_gru(func, opset_version, input_names, output_names,
context):
n_layers, dropout_ratio, hx, ws, bs, xs = func.args
assert n_layers >= 1
hidden_size = hx.shape[2]
gb = onnx_helper.GraphBuilder()
hx_name = input_names[0]
offset = 1
ws_names = [[input_names[offset + i * 6 + j] for j in range(6)]
for i in range(n_layers)]
offset += 6 * n_layers
bs_names = [[input_names[offset + i * 6 + j] for j in range(6)]
for i in range(n_layers)]
offset += 6 * n_layers
xs_names = input_names[offset:]
split_outs = gb.op('Split', [hx_name], num_outputs=n_layers, axis=0)
if n_layers == 1:
split_outs = [split_outs]
# Removing layer dimention and adding num_directions cancels each other
hx_names = split_outs
hy_name, ys_name_list = \
func.reconstruct_return_value(output_names)
y_name = None
hy_names = []
for layer in range(n_layers):
if layer == 0:
# X; shape: (seq_length, batch_size, input_size)
x_name = gb.op(
'Concat',
[gb.op('Unsqueeze', [name], axes=[0]) for name in xs_names],
axis=0)
else:
if opset_version >= 7:
x_name = gb.op('Dropout', [y_name], ratio=dropout_ratio)
elif opset_version >= 6:
x_name = gb.op('Dropout', [y_name], ratio=dropout_ratio,
is_test=0 if chainer.config.train else 1)
else:
x_name = gb.op('Dropout', [y_name], ratio=dropout_ratio,
is_test=0 if chainer.config.train else 1,
consumed_inputs=[1])
# remove num_directions dimention
x_name = gb.op('Squeeze', [x_name], axes=[1])
w = ws_names[layer]
b = bs_names[layer]
# W[zrh]; shape: (num_directions, 3*hidden_size, input_size)
w_name = gb.op(
'Unsqueeze',
[gb.op('Concat', [w[1], w[0], w[2]], axis=0)],
axes=[0])
# R[zrh]; shape: (num_directions, 3*hidden_size, input_size)
r_name = gb.op(
'Unsqueeze',
[gb.op('Concat', [w[4], w[3], w[5]], axis=0)],
axes=[0])
# Wb[zrh], Rb[zrh]; shape: (num_directions, 6*hidden_size)
b_name = gb.op(
'Unsqueeze',
[gb.op('Concat', [b[1], b[0], b[2], b[4], b[3], b[5]], axis=0)],
axes=[0])
# Y; shape: (seq_length, num_directions, batch_size, hidden_size)
# Y_h; shape: (num_directions, batch_size, hidden_size)
y_name, hy_name_ = gb.op(
'GRU',
(x_name, w_name, r_name, b_name, "", hx_names[layer]),
hidden_size=hidden_size,
linear_before_reset=1,
num_outputs=2)
hy_names.append(hy_name_)
split_outs = gb.op(
'Split',
# remove num_directions dimention
[gb.op('Squeeze', [y_name], axes=[1])],
num_outputs=len(ys_name_list), axis=0)
if len(ys_name_list) == 1:
split_outs = [split_outs]
for i, node_name in enumerate(split_outs):
# remove seq_length dimention
gb.op_output_named('Squeeze', [node_name], [ys_name_list[i]], axes=[0])
# Removal of num_directions and new dimention for concatenation cancel
# each other.
gb.op_output_named('Concat', hy_names, [hy_name], axis=0)
return gb.nodes()
def convert_GetItem(func, opset_version, input_names, output_names, context,
parameters):
x = func.inputs[0]
axes, starts, ends = [], [], []
squeeze_idxs, unsqueeze_idxs = [], []
skipped = 0 # when set ellipsis, need to skip index rolling
for i, idx in enumerate(func.slices):
# axis means the index of input x, adjust None and Ellipsis counts
axis = i - len(unsqueeze_idxs) + skipped
if isinstance(idx, slice):
if idx.step is not None and idx.step != 1:
raise ValueError(
'GetItem with {}step slicing is not supported in ONNX '
'Slice operator'.format(idx.step))
axes.append(axis)
starts.append(0 if idx.start is None else idx.start)
ends.append(x.shape[axis] if idx.stop is None else idx.stop)
elif isinstance(idx, int):
axes.append(axis)
starts.append(idx)
ends.append(idx + 1)
squeeze_idxs.append(axis)
elif isinstance(idx, np.ndarray) and idx.ndim == 0:
scalar_idx = idx.item()
axes.append(axis)
starts.append(scalar_idx)
ends.append(scalar_idx + 1)
squeeze_idxs.append(axis)
elif idx is None:
unsqueeze_idxs.append(i - len(squeeze_idxs) + skipped)
elif idx is Ellipsis:
# calculate rest slice number except None, GetItem does not allow
# multiple Ellipsis, so ignore latter Ellipsis count
rest_slice_len = len(
[idx_ for idx_ in func.slices[i + 1:] if idx_ is not None])
assert skipped == 0
skipped = len(x.shape) - axis - rest_slice_len - 1
else:
# not support advanced index like `array[[0,1], [0, 1]]`
raise ValueError(
'GetItem with type {} cannot handle in ONNX Slice, so that '
'ONNX-Chainer does not accept the type'.format(type(idx)))
gb = onnx_helper.GraphBuilder()
if opset_version == 1:
output = gb.op('Slice',
input_names,
axes=axes,
starts=starts,
ends=ends)
elif opset_version == 10:
for param in [('starts', starts), ('ends', ends), ('axes', axes)]:
param_name = context.add_const(
np.asarray(list(param[1]), dtype=np.int64), param[0])
input_names.append(param_name)
output = gb.op('Slice', input_names)
if squeeze_idxs:
output = gb.op('Squeeze', [output], axes=squeeze_idxs)
if unsqueeze_idxs:
output = gb.op('Unsqueeze', [output], axes=unsqueeze_idxs)
return gb.nodes(output_names=output_names)