def template(x_shape=[3, 6],
x_order=OrderNC,
y_order=OrderNC,
indices=AxisKeyDict([Axis.N, Axis.C],
[slice(None), slice(None)]),
description: str = ""):
# vx = np.random.rand(*x_shape)
vx = np.arange(np.product(x_shape)).reshape(x_shape)
x_remain_axes = [a for a in x_order.axes if a in y_order.axes]
vy = np.transpose(
vx[tuple(indices[a] for a in x_order.axes)],
[x_remain_axes.index(a) for a in y_order.axes if a in x_remain_axes])
for i, a in enumerate(y_order.axes):
if a not in x_order.axes:
vy = np.expand_dims(vy, i)
x = Variable(vx.shape, order=x_order)
y, = Slice(None, indices=indices)(x)
y.change_order(y_order)
assert list(vy.shape) == list(y.shape)
generate_kernel_test_case(
description=f"Slice {description}",
graph=Graph([x], [y]),
backend=["webgpu", "webgl", "webassembly"],
inputs={x: vx},
expected={y: vy},
)
def _split_pooling_2d(graph: Graph, op: Pooling2D, v: Variable,
v_pair: Sequence[Variable], axis: Axis):
s1 = v_pair[0].shape_dict[axis]
x = op.inputs["x"]
y = op.outputs["y"]
op.remove_all()
if v == x:
x_0, x_1 = v_pair
s, k, p = (op.SH, op.KH, op.PH) if axis == Axis.H else (op.SW, op.KW,
op.PW)
raise NotImplementedError
elif v == y:
y_0, y_1 = v_pair
s, k, p = (op.SH, op.KH, op.PH) if axis == Axis.H else (op.SW, op.KW,
op.PW)
x_0_range = (0 * s - k // 2, (y_0.shape_dict[axis] - 1) * s + k)
x_1_range = (y_0.shape_dict[axis] * s - k // 2,
(y.shape_dict[axis] - 1) * s + k)
indices = AxisKeyDict(OrderNHWC.axes,
[slice(None) for _ in OrderNHWC.axes])
indices_0 = AxisKeyDict(indices)
indices_0[axis] = slice(max(x_0_range[0], 0),
min(x_0_range[1], x.shape_dict[axis]))
indices_1 = AxisKeyDict(indices)
indices_1[axis] = slice(max(x_1_range[0], 0),
min(x_1_range[1], x.shape_dict[axis]))
x_0, = Slice(None, indices=indices_0)(x)
x_1, = Slice(None, indices=indices_1)(x)
if p > 0:
data = ConstantVariable(
np.zeros([
p if a == axis else x.shape_dict[a] for a in x.order.axes
]), x.order)
x_0, = Concat(None, axis=axis)(data, x_0)
x_1, = Concat(None, axis=axis)(x_1, data)
op_0, op_1 = op.copy(), op.copy()
new_padding = (0, op.PW) if axis == Axis.H else (op.PH, 0)
op_0.parameters["padding"] = new_padding
op_1.parameters["padding"] = new_padding
y_0_new, = op_0(x_0)
y_1_new, = op_1(x_1)
OptimizeRule.replace_variable(graph, y_0_new.transpose_like(y_0), y_0)
OptimizeRule.replace_variable(graph, y_1_new.transpose_like(y_1), y_1)
else:
raise UnexpectedAndPleaseReportError()
def _convert_slice(converter: ONNXConverter, onnx_op: INodeProto):
### TODO Implement this.
x = converter.get_variable(onnx_op.input[0])
attrs = attribute_dict(onnx_op)
print(attrs)
# Attrs:
# 'starts' --> starts.ints
# 'ends' --> ends.ints
# 'axes' --> axes.ints
# https://mil-tokyo.github.io/webdnn/docs/_modules/webdnn/graph/operators/slice.html
# Slice(name, AxisKeyDict: indices)
# eg. multiplier = AxisKeyDict(x.order.axes, [pad_begin if a == axis else x.shape_dict[a] for a in x.order.axes])
# TODO: Construct AxisKeyDict of indices
# Set name for Slice layer if possible?
# Should work like this. I hope.
indices = AxisKeyDict(x.order.axes, [
slice(s, t) for s, t in zip(attrs["starts"].ints, attrs["ends"].ints)
])
#indices = AxisKeyDict(attrs["axes"].ints, [slice(s,t) for s,t in zip(attrs["starts"].ints, attrs["ends"].ints)])
print(x.order.axes)
print(attrs["axes"].ints)
y, = Slice(None, indices)(x)
converter.set_variable(onnx_op.output[0], y)
def slice_handler(converter: TensorFlowConverter, tf_op: "tf.Operation"):
x = converter.get_variable(tf_op.inputs[0])
begin = converter.get_variable(tf_op.inputs[1])
size = converter.get_variable(tf_op.inputs[2])
assert isinstance(
begin, ConstantVariable
), "[TensorFlowConverter] op 'Slice' with dynamic position is not supported yet. "
assert isinstance(
size, ConstantVariable
), "[TensorFlowConverter] op 'Slice' with dynamic size is not supported yet. "
begin = begin.data.flatten().astype(np.int32).tolist()
size = size.data.flatten().astype(np.int32).tolist()
y, = Slice(
None,
indices=AxisKeyDict(x.order.axes,
[slice(b, b + s) for b, s in zip(begin, size)]))(x)
converter.set_variable(tf_op.outputs[0], y)
def _convert_separable_conv2d(converter: KerasConverter,
k_op: "keras.layers.SeparableConv2D"):
x = converter.get_variable(converter.get_input_tensor(k_op)[0])
check_data_format(x, k_op.data_format)
axis_c_in = Axis.C
axis_c_out = Axis()
axis_depth_multiplier = Axis()
w_depthwise = converter.convert_to_constant_variable(
k_op.depthwise_kernel,
Order([Axis.KH, Axis.KW, axis_c_in, axis_depth_multiplier]))
w_pointwise = converter.convert_to_constant_variable(
k_op.pointwise_kernel, Order([Axis.KH, Axis.KW, axis_c_in,
axis_c_out]))
w_pointwise = w_pointwise.reshape(
shape=[
x.shape_dict[axis_c_in], k_op.depth_multiplier,
w_pointwise.shape_dict[axis_c_out]
],
order=Order([axis_c_in, axis_depth_multiplier, axis_c_out]))
ksize = tuple(k_op.kernel_size)
stride = tuple(k_op.strides)
dilation_rate = tuple(k_op.dilation_rate)
padding = (parse_padding(k_op.padding, ksize[0], dilation_rate[0]),
parse_padding(k_op.padding, ksize[1], dilation_rate[1]))
if any(p[0] != p[1] for p in padding):
raise NotImplementedError(
"[KerasConverter] \"Different size padding\" is not supported yet")
padding = tuple(p[0] for p in padding)
h, = Im2Col(None,
ksize=ksize,
stride=stride,
padding=padding,
dilation_rate=dilation_rate)(x)
# TODO: Support depth-wise convolution natively
# Currently, depth-wise convolution is not supported natively, and emulated by composition of small convolution operations.
ys = []
for i in range(h.shape_dict[axis_c_in]):
# 1. Depthwise convolution
#
# Ideal | Current implementation
# ----------------------------------+----------------------------------------------------
# h.axes=[N, H, W, KH, KW, C_in] | g_sub.axes=[N, H, W, KH, KW]
# w.axes=[KH, KW, C_in, DM] | w_sub.axes=[KH, KW, DM]
# g.axes=[N, H, W, C_in, DM] | g_sub.axes=[N, H, W, DM]
h_sub, = Slice(
None,
indices=AxisKeyDict(
h.order.axes,
[i if a == axis_c_in else slice(None)
for a in h.order.axes]))(h)
w_depthwise_sub = w_depthwise[:, :, i, :]
g_sub, = Tensordot(None, axes=((Axis.KH, Axis.KW),
(Axis.KH, Axis.KW)))(h_sub,
w_depthwise_sub)
# 2. Pointwise (projection) convolution
#
# Ideal | Current implementation
# ----------------------------------+----------------------------------------------------
# g.axes=[N, H, W, C_in, DM] | g_sub.axes=[N, H, W, DM]
# w.axes=[DM, Cin, C_out] | w_sub.axes=[DM, C_out]
# y.axes=[N, H, W, C_out] | y_sub.axes=[N, H, W, C_out]
w_pointwise_sub = w_pointwise[i, :, :]
y_sub, = Tensordot(None,
axes=((axis_depth_multiplier, ),
(axis_depth_multiplier, )))(g_sub,
w_pointwise_sub)
ys.append(y_sub)
# Sum up all sub convolution results to one
while len(ys) > 1:
ys.append(ys.pop(0) + ys.pop(0))
y = ys[0]
# reinterpret axis "C_out" as C
axes = list(y.order.axes)
i = axes.index(axis_c_out)
axes.pop(i)
axes.insert(i, Axis.C)
y = y.reinterpret_axes(Order(axes))
if k_op.data_format == "channels_last":
y = y.transpose(OrderNHWC)
elif k_op.data_format == "channels_first":
y = y.transpose(OrderNCHW)
else:
raise NotImplementedError(
f"[KerasConverter] Unknown data format: {k_op.data_format}")
if k_op.use_bias:
b = converter.convert_to_constant_variable(k_op.bias, OrderC)
y = y + b
y = do_activation(k_op.activation, y)
converter.set_variable(converter.get_output_tensor(k_op)[0], y)