This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 7903ae1451 [Unity][Frontends][Onnx] Improve ConstantOfShape behavior
(#15434)
7903ae1451 is described below
commit 7903ae1451cbe144e426b305e1232a6768fc650f
Author: Josh Fromm <jwfr...@octoml.ai>
AuthorDate: Sat Jul 29 01:24:21 2023 -0700
[Unity][Frontends][Onnx] Improve ConstantOfShape behavior (#15434)
This small PR improves how the onnx importer deals with ConstantOfShape.
Previously we were creating constants that could be massive if the provided
shape was large. This caused an explosion in the resulting model parameter
size. Using `broadcast_to` instead provides behavior closer to the intention of
the source graph.
I also noticed relax now has `relax.op.image.resize2d` so i replaced our
topi call with it which should help produce more readable output graphs.
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 52 +++++++++----------------
1 file changed, 19 insertions(+), 33 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 9ec340c038..9711ee2578 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -249,7 +249,6 @@ class Unsqueeze(OnnxOpConverter):
def _impl_v13(cls, bb, inputs, attr, params):
data = inputs[0]
axes = get_constant(inputs[1], params)
-
# If input is a constant, compute directly
if isinstance(data, relax.Constant) and isinstance(axes,
relax.Constant):
axes = axes.data.numpy().tolist()
@@ -666,34 +665,20 @@ class ConstantOfShape(OnnxOpConverter):
else:
dtype = "float32"
- # If shape is a constant, we can directly create a relax constant.
+ # If shape is a constant, treat it as a shapexpr.
if isinstance(shape, relax.Constant):
- np_array = _np.zeros(shape=shape.data.numpy()) + value
- return relax.const(np_array, dtype=dtype)
- elif isinstance(shape, relax.ShapeExpr):
- np_array = _np.zeros(shape=[dim.value for dim in shape]) + value
- return relax.const(np_array, dtype)
+ shape = relax.ShapeExpr(list(shape.data.numpy()))
- # Otherwise we have to use the value of shape at runtime.
# Create a constant for the new value.
const_value = relax.const(value, dtype)
- # Convert to shape expression if needed.
- if not isinstance(shape.struct_info, relax.ShapeStructInfo):
- shape_ndim = [dim.value for dim in
shape.struct_info.shape.values][0]
- # Broadcast the constant to the input shape.
- shape_dataflow_var = bb.emit(
- relax.Call(
- relax.ExternFunc("vm.builtin.tensor_to_shape"),
- [shape],
- sinfo_args=[relax.ShapeStructInfo(ndim=shape_ndim)],
- )
- )
- shape_vars = []
- for i in range(shape_ndim):
- shape_vars.append(tvm.tir.Var("x_%d" % i, "int64"))
- bb.match_cast(shape_dataflow_var,
relax.ShapeStructInfo(shape_vars))
- shape = relax.ShapeExpr(shape_vars)
+ # Special case where requested shape is a scalar.
+ if len(shape) == 1 and int(shape[0]) == 1:
+ return const_value
+
+ # Convert to shape expression from tensor if needed.
+ if not isinstance(shape, relax.ShapeExpr):
+ shape = relax.op.tensor_to_shape(shape)
return relax.op.broadcast_to(const_value, shape)
@@ -983,8 +968,9 @@ class Expand(OnnxOpConverter):
# For some reason, onnx allows target shapes to be smaller than
input shapes.
# We need to go correct it.
data_shape = [dim.value for dim in data.struct_info.shape]
+ # Fix small target shapes.
for i, s in enumerate(new_shape):
- if s < data_shape[i]:
+ if i < len(data_shape) and s < data_shape[i]:
new_shape[i] = data_shape[i]
# If the new shape matches the input shape, no transformation is
needed.
if new_shape == data_shape:
@@ -1200,18 +1186,16 @@ class Resize(OnnxOpConverter):
), "Only constant output size currently supported."
sizes = sizes.data.numpy().astype("int64").tolist()[2:]
- # TODO(jwfromm) relax.image.resize2d runs into some issues with
dynamism.
- return bb.emit_te(
- topi.image.resize2d,
+ return relax.op.image.resize2d(
x,
- roi,
- sizes,
+ size=relax.ShapeExpr(sizes),
+ roi=roi,
layout="NCHW",
method=mode,
coordinate_transformation_mode=coord_mode,
rounding_method=rounding_method,
- bicubic_alpha=cubic_coeff_a,
- bicubic_exclude=exclude_outside,
+ cubic_alpha=cubic_coeff_a,
+ cubic_exclude=exclude_outside,
extrapolation_value=extrapolation_value,
)
@@ -1880,7 +1864,9 @@ class ONNXGraphImporter:
array = self._parse_array(init_tensor)
# Create variables for constants.
if self._keep_params_in_input:
- init_var = self._new_var(init_tensor.name, shape=array.shape,
dtype=array.dtype)
+ # Pytorch sometimes inserts silly weight prefix. Remove it.
+ var_name = init_tensor.name.strip("onnx::")
+ init_var = self._new_var(var_name, shape=array.shape,
dtype=array.dtype)
self._nodes[init_tensor.name] = init_var
# We need to keep track of both the real value and variable
for this variable.
self._params[init_tensor.name] = (init_var, array)
