def _optimize_redundant_transposed_output(
graph: Graph,
op: Operator,
var_name: str,
target_orders: Optional[Union[Order, List[Order]]] = None):
v = op.outputs[var_name]
if len(v.input_to) != 1:
return False
op2 = list(v.input_to)[0]
if not isinstance(op2, Transpose):
return False
if target_orders is not None:
if isinstance(target_orders, Order):
target_orders = [target_orders]
if op2.outputs["y"].order not in target_orders:
return False
v2 = op2.outputs["y"]
op2.remove_all()
OptimizeRule.replace_variable(graph, v, v2, with_assert=False)
return True
def fold_constance(self, graph: Graph):
x0 = self.inputs["x0"] # type: ConstantVariable
y = self.outputs["y"]
self.remove_all()
new_y = ConstantVariable(x0.copy().change_order(y.order).data ** self.value, y.order)
OptimizeRule.replace_variable(graph, y, new_y)
def _split_im2col(graph: Graph, op: Im2Col, v: Variable, v_pair: Sequence[Variable], axis: Axis):
s1 = v_pair[0].shape_dict[axis]
im = op.inputs["im"]
col = op.outputs["col"]
op.remove_all()
if v == col:
"""
before)
im -{Im2Col}- col
after)
+- col_0
im -{PartialIm2Col}-+
+- col_1
"""
col_0, col_1 = PartialIm2Col(None,
ksize=op.ksize, stride=op.stride, padding=op.padding, dilation_rate=op.dilation_rate,
axis=axis, sections=[s1])(im)
OptimizeRule.replace_variable(graph, col_0.transpose(v_pair[0].order), v_pair[0])
OptimizeRule.replace_variable(graph, col_1.transpose(v_pair[1].order), v_pair[1])
elif v == im:
raise NotImplementedError(f"Variable is too large to handle in WebGL backend: {v}")
else:
raise UnexpectedAndPleaseReportError
def optimize(self, graph: Graph) -> Tuple[Graph, bool]:
flag_changed = False
for op in traverse.filter_nodes(
traverse.listup_operators(graph),
Deconvolution2D): # type: Deconvolution2D
x = op.inputs["x"]
w = op.inputs["w"]
y = op.outputs["y"]
flag_changed = True
op.remove_all()
a_filter, a_kh, a_kw = Axis(), Axis(), Axis()
w, = ReinterpretAxis(None,
in_order=OrderNHWC,
out_order=Order(
[Axis.C, a_kh, a_kw, a_filter]))(w)
x, = ReinterpretAxis(None,
in_order=OrderNHWC,
out_order=Order(
[Axis.N, Axis.H, Axis.W, a_filter]))(x)
col, = Tensordot(None, axes=a_filter)(x, w)
col = col.transpose(
Order([Axis.N, Axis.H, Axis.W, a_kh, a_kw, Axis.C]))
col = col.reshape(shape=[*col.shape[0:3],
mul(col.shape[3:6])],
order=OrderNHWC)
new_y, = Col2Im(None,
ksize=op.ksize,
stride=op.stride,
padding=op.padding)(col)
OptimizeRule.replace_variable(graph, new_y.transpose_like(y), y)
return graph, flag_changed
def fold_constance(self, graph: Graph):
x0 = self.inputs["x0"]
y = self.outputs["y"]
OptimizeRule.replace_variable(graph, y,
x0.copy().change_order(y.order))
self.remove_all()
def optimize(self, graph: Graph):
flag_changed = False
for x in graph.inputs:
if len(x.input_to) != 1:
continue
op = list(x.input_to)[0]
if isinstance(op, ConvertRGBAtoR) or isinstance(
op, ConvertRtoRGBA):
flag_changed = True
y = op.outputs["y"]
op.remove_all()
OptimizeRule.replace_variable(graph, x, y)
for y in graph.outputs:
if isinstance(y.output_from, ConvertRGBAtoR) or isinstance(
y.output_from, ConvertRtoRGBA):
flag_changed = True
x = y.output_from.inputs["x0"]
i = graph.outputs.index(y)
graph.outputs.remove(y)
graph.outputs.insert(i, x)
if len(y.input_to) == 0:
y.output_from.remove_all()
return graph, flag_changed
def optimize(self, graph: Graph) -> Tuple[Graph, bool]:
flag_changed = False
for op in traverse.filter_nodes(traverse.listup_operators(graph),
Convolution2D): # type: Convolution2D
x = op.inputs["x"]
w = op.inputs["w"]
y = op.outputs["y"]
flag_changed = True
op.remove_all()
a_filter, a_kh, a_kw = Axis(), Axis(), Axis()
w, = ReinterpretAxis(None,
in_order=OrderNHWC,
out_order=Order(
[Axis.C, a_kh, a_kw, a_filter]))(w)
if op.WH == 1 and op.WW == 1 and op.stride == (
1, 1) and op.padding == (0, 0):
# Projection
col, = ReinterpretAxis(
None,
in_order=OrderNHWC,
out_order=Order([Axis.N, Axis.H, Axis.W, a_filter]))(x)
new_y, = Tensordot(None,
[[a_filter], [a_kh, a_kw, a_filter]])(col,
w)
elif op.WH == x.shape_dict[Axis.H] and op.WW == x.shape_dict[
Axis.W] and op.padding == (0, 0):
# Global convolution
col, = ReinterpretAxis(None,
in_order=OrderNHWC,
out_order=Order(
[Axis.N, a_kh, a_kw, a_filter]))(x)
new_y, = Tensordot(
None, [[[a_kh, a_kw, a_filter], [a_kh, a_kw, a_filter]],
[a_kh, a_kw, a_filter]])(col, w)
else:
# General convolution
col, = Im2Col(None,
ksize=op.ksize,
stride=op.stride,
padding=op.padding,
dilation_rate=op.dilation_rate)(x)
col, = ReinterpretAxis(
None,
in_order=OrderNHWC,
out_order=Order([Axis.N, Axis.H, Axis.W, a_filter]))(col)
new_y, = Tensordot(None,
[[a_filter], [a_kh, a_kw, a_filter]])(col,
w)
new_y = new_y.transpose(y.order)
OptimizeRule.replace_variable(graph, new_y, y)
return graph, flag_changed
def optimize(self, graph: Graph) -> Tuple[Graph, bool]:
flag_changed = False
for concat in traverse.filter_nodes(traverse.listup_operators(graph), Concat):
if len(concat.inputs) == 2:
# Unrolling is not needed
continue
flag_changed = True
xs = [concat.inputs[f"x{i}"] for i in range(len(concat.inputs))]
y = concat.outputs["y"]
concat.remove_all()
while len(xs) > 1:
hs = []
while len(xs) > 0:
if len(xs) == 1:
hs.append(xs.pop(0))
else:
x0, x1 = xs.pop(0), xs.pop(0)
h, = Concat(None, axis=concat.axis)(x0, x1)
hs.append(h)
xs = hs
OptimizeRule.replace_variable(graph, y, xs[0].transpose_like(y))
return graph, flag_changed
def fold_constance(self, graph: "graph.Graph"):
x = self.inputs["x"] # type: ConstantVariable
y = self.outputs["y"]
self.remove_all()
new_y = ConstantVariable(x.data.copy(), y.order)
OptimizeRule.replace_variable(graph, y, new_y)
def optimize(self, graph: Graph) -> Tuple[Graph, bool]:
flag_changed = False
for op in traverse.filter_nodes(traverse.listup_operators(graph),
Linear):
x = op.inputs["x"]
w = op.inputs["w"]
y = op.outputs["y"]
flag_changed = True
op.remove_all()
a_filter = Axis()
if x.ndim == 2:
w, = ReinterpretAxis(None,
in_order=OrderNC,
out_order=Order([Axis.C, a_filter]))(w)
new_y, = Tensordot(None, axes=[Axis.C, a_filter])(x, w)
elif x.ndim == 4:
w, = ReinterpretAxis(
None,
in_order=OrderNHWC,
out_order=Order([Axis.C, Axis.H, Axis.W, a_filter]))(w)
new_y, = Tensordot(None,
axes=[[Axis.H, Axis.W, Axis.C],
[Axis.H, Axis.W, a_filter]])(x, w)
else:
raise NotImplementedError
OptimizeRule.replace_variable(graph, new_y.transpose_like(y), y)
return graph, flag_changed
def fold_constance(self, graph: Graph):
x = self.inputs["x"] # type: ConstantVariable
y = self.outputs["y"]
new_y = ConstantVariable(np.tile(x.data, self.multiplier), x.order)
new_y.change_order(y.order)
OptimizeRule.replace_variable(graph, y, new_y)
self.remove_all()
def fold_constance(self, graph: Graph):
xs = [self.inputs[f"x{i}"] for i in range(len(self.inputs))] # type: List[ConstantVariable]
y = self.outputs["y"]
data = np.concatenate([x.copy().change_order(y.order).data for x in xs], axis=y.order.axes_dict[self.axis])
new_y = ConstantVariable(data, y.order)
OptimizeRule.replace_variable(graph, y, new_y)
self.remove_all()
def fold_constance(self, graph: Graph):
x = self.inputs["x"] # type: ConstantVariable
y = self.outputs["y"]
self.remove_all()
y_new = ConstantVariable(x.data, x.order).change_order(self.in_order)
y_new = ConstantVariable(y_new.data.reshape(self.out_shape), self.out_order).change_order(y.order)
OptimizeRule.replace_variable(graph, y, y_new)
def fold_constance(self, graph: Graph):
x0 = self.inputs["x0"] # type: ConstantVariable
y = self.outputs["y"]
self.remove_all()
y_new = ConstantVariable(x0.data, x0.order).change_order(y.order)
y_new.data = y_new.data**self.value
OptimizeRule.replace_variable(graph, y, y_new)
def fold_constance(self, graph: Graph):
x0 = self.inputs["x0"] # type: ConstantVariable
y = self.outputs["y"]
new_y = ConstantVariable(1 / np.sqrt(x0.data), x0.order)
new_y.change_order(y.order)
OptimizeRule.replace_variable(graph, y, new_y)
self.remove_all()
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 _split_partial_im2col(graph: Graph, op: PartialIm2Col, v: Variable,
v_pair: Sequence[Variable], axis: Axis):
s1 = v_pair[0].shape_dict[axis]
im = op.inputs["im"]
cols = [op.outputs[f"col{i}"] for i in range(len(op.outputs))]
sections = op.sections
if v == im:
raise NotImplementedError(
f"Variable is too large to handle in WebGL backend: {v}")
elif v in cols:
op.remove_all()
if axis == op.axis:
"""
before)
+- col0
|
im -{PartialIm2Col}-+- col1
|
+- col2
after)
+- col0
|
+- col1_0
im -{PartialIm2Col}-+
+- col1_1
|
+- col2
"""
target_i = cols.index(v)
s_insert = (0 if target_i == 0 else sections[target_i - 1]) + s1
new_sections = list(sections)
new_sections.insert(target_i, s_insert)
cols.pop(target_i)
cols.insert(target_i + 0, v_pair[0])
cols.insert(target_i + 1, v_pair[1])
new_cols = PartialIm2Col(None,
ksize=op.ksize,
stride=op.stride,
padding=op.padding,
dilation_rate=op.dilation_rate,
axis=axis,
sections=new_sections)(im)
for col, new_col in zip(cols, new_cols):
OptimizeRule.replace_variable(graph, new_col, col)
else:
raise NotImplementedError(
f"Variable is too large to handle in WebGL backend: {v}")
else:
raise UnexpectedAndPleaseReportError
def optimize(self, graph: Graph):
flag_changed = False
"""
before)
v0[RGBA] -{ConvertRtoRGBA}- v1[RGBA]
after)
v0[RGBA] -{ConvertRGBAtoR}- v2[Order=v0.order][R] -{Transpose}- v3[Order=v1.order][R]-{ConvertRtoRGBA}- v1[RGBA]
"""
matches = traverse.search_sub_structure(
graph, [Variable, ConvertRtoRGBA, Variable])
while len(matches) > 0:
v0, r2rgba, v1 = matches.pop(
) # type: Variable, ConvertRtoRGBA, Variable
if not (ChannelMode.get(v0) == ChannelMode.get(v1) ==
ChannelModeEnum.RGBA):
continue
flag_changed = True
r2rgba.remove_all()
v2 = convert_rgba_to_r(v0)
v2.change_order(v0.order)
v3 = v2.transpose(v1.order)
v1_new = convert_r_to_rgba(v3)
v1_new.change_order(v1.order)
OptimizeRule.replace_variable(graph, v1_new, v1)
"""
before)
v0[R] -{ConvertRGBAtoR}- v1[R]
after)
v0[R] -{Transpose}- v1[R]
"""
matches = traverse.search_sub_structure(
graph, [Variable, ConvertRGBAtoR, Variable])
while len(matches) > 0:
v0, rgba2r, v1 = matches.pop(
) # type: Variable, ConvertRGBAtoR, Variable
if not (ChannelMode.get(v0) == ChannelMode.get(v1) ==
ChannelModeEnum.R):
continue
flag_changed = True
rgba2r.remove_all()
OptimizeRule.replace_variable(graph, v0.transpose(v1.order), v1)
return graph, flag_changed
def optimize(self, graph: Graph) -> Tuple[Graph, bool]:
flag_changed = False
variables = traverse.listup_variables(graph)
while len(variables) > 0:
x = variables.pop()
for op1, op2 in itertools.permutations(x.input_to, 2):
if op2 is op1:
continue
if op2.__class__ != op1.__class__:
# class is not same
continue
if any((x_name not in op2.inputs) or (
op2.inputs[x_name] != op1.inputs[x_name])
for x_name in op1.inputs.keys()):
# input is not same
continue
if any((key not in op2.parameters) or (
op2.parameters[key] != op1.parameters[key])
for key in op1.parameters.keys()):
# parameter is not same
continue
flag_changed = True
vs_1 = dict(op1.outputs)
vs_2 = dict(op2.outputs)
op2.remove_all()
for v_name, v1 in vs_1.items():
v2 = vs_2[v_name]
if v1.order == v2.order:
"""
+-{op3}-
-{op1}- v1 -+
+-{op4}-
"""
OptimizeRule.replace_variable(graph, v2, v1)
else:
"""
+-{op3}-
-{op1}- v1 -+
+-{Transpose}- v2 -{op4}-
"""
v2_dummy, = Transpose(None)(v1)
v2_dummy.change_order(v2.order)
OptimizeRule.replace_variable(graph, v2_dummy, v2)
variables = traverse.listup_variables(graph)
break
return graph, flag_changed
def fold_constance(self, graph: Graph):
x0 = self.inputs["x0"] # type: ConstantVariable
y = self.outputs["y"]
self.remove_all()
OptimizeRule.replace_variable(
graph, y,
ConstantVariable(np.arccosh(x0.copy().change_order(y.order).data),
y.order))
def fold_constance(self, graph: Graph):
in_order = self.parameters["in_order"]
out_shape = self.parameters["out_shape"]
out_order = self.parameters["out_order"]
x = self.inputs["x"] # type: ConstantVariable
y = self.outputs["y"]
self.remove_all()
new_y = ConstantVariable(x.copy().change_order(in_order).data.reshape(out_shape), out_order)
OptimizeRule.replace_variable(graph, y, new_y)
def fold_constance(self, graph: "graph.Graph"):
x = self.inputs["x"] # type: ConstantVariable
y = self.outputs["y"]
self.remove_all()
y_new = ConstantVariable(
x.data,
Order([
self.out_order.axes[self.in_order.axes.index(a)]
for a in x.order.axes
]))
OptimizeRule.replace_variable(graph, y, y_new.change_order(y.order))
def fold_constance(self, graph: Graph):
x = self.inputs["x"] # type: ConstantVariable
y = self.outputs["y"]
remained_axes_in_x_order = [a for a in x.order.axes if a in y.order.axes]
new_axes = [a for a in y.order.axes if a not in x.order.axes]
slices = [self.indices[a] for a in x.order.axes] + [None] * len(new_axes)
new_y = ConstantVariable(x.data[slices], Order(remained_axes_in_x_order + new_axes))
new_y.change_order(y.order)
OptimizeRule.replace_variable(graph, y, new_y)
self.remove_all()
def fold_constance(self, graph: Graph):
x = self.inputs["x"] # type: ConstantVariable
ys = [self.outputs[f"y{i}"] for i in range(len(self.outputs))]
axis = self.parameters["axis"]
sections = self.parameters["sections"]
self.remove_all()
y_datum = np.split(x.data, sections, x.order.axes_dict[axis])
for i, y in enumerate(ys):
y_new = ConstantVariable(y_datum[i], x.order).change_order(y.order)
OptimizeRule.replace_variable(graph, y, y_new)
def fold_constance(self, graph: Graph):
x = self.inputs["x"] # type: ConstantVariable
y = self.outputs["y"]
new_axes = list(x.order.axes)
new_axes.remove(self.axis)
new_y = ConstantVariable(
np.sum(x.data, axis=x.order.axes_dict[self.axis]), Order(new_axes))
new_y.change_order(y.order)
OptimizeRule.replace_variable(graph, y, new_y)
self.remove_all()
def optimize(self, graph: Graph) -> Tuple[Graph, bool]:
flag_changed = False
matches = traverse.search_sub_structure(
graph, [Tensordot, Variable, ElementwiseMul, Variable])
while len(matches) > 0:
tensordot, h, elementwise_mul, y = matches.pop(
) # type: Tensordot, Variable, ElementwiseMul, Variable
if len(h.input_to) != 1:
# h will be removed by this optimization rule
continue
if isinstance(tensordot.inputs["A"], ConstantVariable):
w1 = tensordot.inputs["A"]
reduced_axes = tensordot.axes[0]
elif isinstance(tensordot.inputs["B"], ConstantVariable):
w1 = tensordot.inputs["B"]
reduced_axes = tensordot.axes[1]
else:
continue
if isinstance(
elementwise_mul.inputs["x0"],
ConstantVariable) and elementwise_mul.inputs["x1"] == h:
w2 = elementwise_mul.inputs["x0"]
elif isinstance(
elementwise_mul.inputs["x1"],
ConstantVariable) and elementwise_mul.inputs["x0"] == h:
w2 = elementwise_mul.inputs["x1"]
else:
continue
if any(axis not in w1.order.axes for axis in w2.order.axes):
continue
if any(axis in reduced_axes for axis in w2.order.axes):
continue
flag_changed = True
elementwise_mul.remove_all()
OptimizeRule.replace_variable(graph,
w1,
w1.copy() * w2,
with_assert=False)
OptimizeRule.replace_variable(graph, h, y, with_assert=False)
return graph, flag_changed
def optimize_operator(self, graph: Graph, op: ReinterpretAxis):
x = op.inputs["x"]
y = op.outputs["y"]
if len(x.input_to) == 1 and x.output_from is None:
op.remove_all()
if isinstance(x, ConstantVariable):
x = ConstantVariable(x.data, y.order)
if y in graph.outputs:
index = graph.outputs.index(y)
graph.outputs.remove(y)
graph.outputs.insert(index, x)
else:
OptimizeRule.replace_variable(graph, y, x)
else:
assert x in graph.inputs
index = graph.inputs.index(x)
graph.inputs.remove(x)
graph.inputs.insert(index, y)
return True
if op.parameters["in_order"] == op.parameters["out_order"]:
_remove_unary_operator(graph, op)
return True
flag_changed = False
for axis1, axis2 in zip(op.parameters["in_order"].axes,
op.parameters["out_order"].axes):
is_resolved1 = not (isinstance(axis1, AxisVar)
and axis1.value is None)
is_resolved2 = not (isinstance(axis2, AxisVar)
and axis2.value is None)
if is_resolved1 and not is_resolved2:
axis2.unify(axis1)
flag_changed = True
elif not is_resolved1 and is_resolved2:
axis1.unify(axis2)
flag_changed = True
if flag_changed:
return True
return False
def optimize_pair(self, graph: Graph, op1: ElementwiseAdd, op2: ElementwiseDiv):
c1, v1 = _get_constant_and_variable(op1, "x0", "x1")
if c1 is None:
return False
c2, v2 = _get_constant_and_variable(op2, "x0", "x1")
if c2 is None:
return False
y2 = op2.outputs["y"]
op2.remove_all()
op1.remove_all()
y = (v1 / c2) + (c1 / c2)
OptimizeRule.replace_variable(graph, y2, y.change_order(y2.order))
return True
def _split_tensorwise(graph: Graph, op: Operator, v: Variable,
v_pair: Sequence[Variable], axis: Axis):
s1 = v_pair[0].shape_dict[axis]
s2 = v_pair[1].shape_dict[axis]
xs = dict(op.inputs)
ys = dict(op.outputs)
op.remove_all()
op_0 = op.copy()
op_1 = op.copy()
for key, x in xs.items():
if x == v:
x_0, x_1 = v_pair
else:
if axis in x.order.axes:
x_0, x_1 = SplitAxis(None, axis=axis, sections=[s1])(x)
else:
# splitting is not occurred
x_0 = x_1 = x
op_0.append_input(key, x_0)
op_1.append_input(key, x_1)
for key, y in ys.items():
if y == v:
y_0, y_1 = v_pair
else:
if axis in y.order.axes:
# TODO (Kiikurage)
# Attribute attached to "y" is not copied to neither "y_0" or "y_1"
y_0 = Variable([
s1 if a == axis else y.shape_dict[a] for a in y.order.axes
], y.order)
y_1 = Variable([
s2 if a == axis else y.shape_dict[a] for a in y.order.axes
], y.order)
y_new, = Concat(None, axis=axis)(y_0, y_1)
OptimizeRule.replace_variable(graph, y, y_new)
else:
raise UnexpectedAndPleaseReportError
op_0.append_output(key, y_0)
op_1.append_output(key, y_1)
def optimize_operator(self, graph: Graph, op: Reshape):
x = op.inputs["x"]
y = op.outputs["y"]
if x.order == y.order and x.shape == y.shape:
# no reshape is required
_remove_unary_operator(graph, op)
return True
if x.shape == y.shape:
# only reinterpret_axis is required
op.remove_all()
y_dummy = x.reinterpret_axes(y.order)
OptimizeRule.replace_variable(graph, y_dummy, y)
return True
return False
