def verify(dshape, begin, end, strides, output, slice_mode="end", test_ref=True, dtype="int32"): x = relay.var("x", relay.TensorType(dshape, "float32")) ndim = len(dshape) begin = begin if begin else [0] * ndim end = end if end else list(dshape) if strides: if len(strides) == 1: strides = strides * ndim else: strides = [1] * ndim # target numpy result x_data = np.random.uniform(size=dshape).astype("float32") ref_res = tvm.topi.testing.strided_slice_python(x_data, begin, end, strides, slice_mode) data = [x_data, np.array(begin), np.array(end)] begin = relay.const(begin, dtype=dtype) end = relay.const(end, dtype=dtype) if strides: data.append(np.array(strides)) strides = relay.const(strides, dtype=dtype) z = relay.strided_slice(x, begin=begin, end=end, strides=strides, slice_mode=slice_mode) else: z = relay.strided_slice(x, begin=begin, end=end, slice_mode=slice_mode) func = relay.Function([x], z) func = run_infer_type(func) func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType()) assert isinstance(func2.body, relay.Call) assert func2.body.op == relay.op.get("strided_slice") verify_func(func2, [x_data], ref_res)
def expected(x, w1, w2, b1, b2, j, bias_shape1, bias_shape2): args = [x, w1, w2, b1, b2] w_stacked = relay.concatenate((w1, w2), axis=0) y = relay.nn.dense(x, w_stacked, units=3 * j) n_out_dims = max(len(bias_shape1), 2) if len(bias_shape1) == 0: b1 = relay.repeat(relay.expand_dims(b1, -1), j, 0) elif bias_shape1[-1] == 1: b1 = relay.repeat(b1, j, len(bias_shape1) - 1) if len(bias_shape2) == 0: b2 = relay.repeat(relay.expand_dims(b2, -1), 2 * j, 0) elif bias_shape2[-1] == 1: b2 = relay.repeat(b2, 2 * j, len(bias_shape2) - 1) b = relay.concatenate((b1, b2), axis=max(0, len(bias_shape1) - 1)) y = relay.add(y, b) begin = [0 for _ in range(n_out_dims - 1)] end = [-1 for _ in range(n_out_dims - 1)] strides = [1 for _ in range(n_out_dims)] y1 = relay.strided_slice(y, begin=begin + [0], end=end + [j], strides=strides, slice_mode="size") y2 = relay.strided_slice(y, begin=begin + [j], end=end + [2 * j], strides=strides, slice_mode="size") return relay.Function(args, relay.Tuple((y1, y2)))
def expected(x, w1, w2, w3, w4, channels1, channels2, channels3, channels4): # use a fixed order of args so alpha equal check can pass args = [x, w1, w2, w3, w4] w = relay.concatenate((w1, w2, w4), axis=0) y = relay.nn.conv2d(x, w, channels=channels1 + channels2 + channels4) y1 = relay.strided_slice(y, begin=relay.const([0, 0], "int64"), end=relay.const([-1, channels1], "int64"), strides=relay.const([1, 1], 'int64'), slice_mode="size") y2 = relay.strided_slice(y, begin=relay.const([0, channels1], "int64"), end=relay.const([-1, channels2], "int64"), strides=relay.const([1, 1], 'int64'), slice_mode="size") y3 = relay.nn.conv2d(x, w3) y4 = relay.strided_slice(y, begin=relay.const([0, channels1 + channels2], "int64"), end=relay.const([-1, channels4], "int64"), strides=relay.const([1, 1], 'int64'), slice_mode="size") y5 = relay.nn.max_pool2d(x) y = relay.Tuple((y1, y2, y3, y4, y5)) return relay.Function(args, y)
def expected(x, w1, w2, b1, b2, scale1, scale2, newshape1, newshape2, j): args = [x, w1, w2, b1, b2, scale1, scale2] w_stacked = relay.concatenate((w1, w2), axis=0) y = relay.nn.dense(x, w_stacked, units=3 * j) b = relay.concatenate((b1, b2), axis=0) y = relay.add(y, b) scale1 = relay.repeat(scale1, j, 0) scale2 = relay.repeat(scale2, 2 * j, 0) scale = relay.concatenate((scale1, scale2), axis=0) y = relay.multiply(y, scale) strides = [1, 1] y1 = relay.strided_slice(y, begin=[0, 0], end=[-1, j], strides=strides, slice_mode="size") y2 = relay.strided_slice(y, begin=[0, j], end=[-1, 2 * j], strides=strides, slice_mode="size") y1 = relay.reshape(y1, newshape=newshape1) y2 = relay.reshape(y2, newshape=newshape2) y = relay.Tuple((y1, y2)) return relay.Function(args, y)
def verify(dshape, begin, end, strides, output, slice_mode="end", test_ref=True, dtype="int32"): x = relay.var("x", relay.TensorType(dshape, "float32")) ndim = len(dshape) begin = begin if begin else [0] * ndim end = end if end else list(dshape) if strides: if len(strides) == 1: strides = strides * ndim else: strides = [1] * ndim # target numpy result x_data = np.random.uniform(size=dshape).astype("float32") ref_res = tvm.topi.testing.strided_slice_python( x_data, begin, end, strides, slice_mode) data = [ x_data, np.array(begin, dtype=dtype), np.array(end, dtype=dtype) ] begin = relay.var("begin", shape=[len(begin)], dtype=dtype) end = relay.var("end", shape=[len(end)], dtype=dtype) inputs = [x, begin, end] if strides: data.append(np.array(strides, dtype=dtype)) strides = relay.var("strides", shape=[len(strides)], dtype=dtype) inputs.append(strides) z = relay.strided_slice(x, begin=begin, end=end, strides=strides, slice_mode=slice_mode) else: z = relay.strided_slice(x, begin=begin, end=end, slice_mode=slice_mode) func = relay.Function(inputs, z) func = run_infer_type(func) text = func.astext() if not test_ref: return for target, dev in tvm.testing.enabled_targets(): mod = tvm.ir.IRModule.from_expr(func) intrp = relay.create_executor("vm", mod=mod, device=dev, target=target) op_res = intrp.evaluate()(*data) tvm.testing.assert_allclose(op_res.asnumpy(), ref_res)
def verify(dshape, begin, end, strides, output, slice_mode="end", test_ref=True, dtype="int32"): ndim = len(dshape) begin = begin if begin else [0] * ndim end = end if end else list(dshape) # target numpy result x_data = np.random.uniform(size=dshape).astype("float32") ref_res = tvm.topi.testing.strided_slice_python(x_data, begin, end, strides, slice_mode) x = relay.var("x", relay.TensorType((relay.Any(),) * ndim, "float32")) if strides: z = relay.strided_slice(x, begin=begin, end=end, strides=strides, slice_mode=slice_mode) else: z = relay.strided_slice(x, begin=begin, end=end, slice_mode=slice_mode) func = relay.Function([x], z) func = run_infer_type(func) text = func.astext() assert "begin=" in text assert "end=" in text if not test_ref: return for target, ctx in tvm.testing.enabled_targets(): mod = tvm.ir.IRModule.from_expr(func) intrp = relay.create_executor("vm", mod=mod, ctx=ctx, target=target) op_res = intrp.evaluate()(x_data) tvm.testing.assert_allclose(op_res.asnumpy(), ref_res)
def expected(x, w1, w2, w3, b1, b2, b3): # use a fixed order of args so alpha equal check can pass s1 = w1.type_annotation.shape[1] s2 = w2.type_annotation.shape[1] s3 = w3.type_annotation.shape[1] args = [x, w1, w2, w3, b1, b2, b3] w = relay.concatenate((w1, w2, w3), axis=1) b = relay.concatenate((b1, b2, b3), axis=-1) y = relay.nn.batch_matmul(x, w) y = relay.add(y, b) y1 = relay.strided_slice(y, begin=relay.const([0, 0, 0], "int64"), end=relay.const([-1, -1, s1], "int64"), strides=relay.const([1, 1, 1], 'int64'), slice_mode="size") y2 = relay.strided_slice(y, begin=relay.const([0, 0, s1], "int64"), end=relay.const([-1, -1, s2], "int64"), strides=relay.const([1, 1, 1], 'int64'), slice_mode="size") y3 = relay.strided_slice(y, begin=relay.const([0, 0, s1+s2], "int64"), end=relay.const([-1, -1, s3], "int64"), strides=relay.const([1, 1, 1], 'int64'), slice_mode="size") y = relay.Tuple((y1, y2, y3)) return relay.Function(args, y)
def verify( dshape, begin, end, strides, output, axes=None, slice_mode="end", test_ref=True, dtype="int32", ): x = relay.var("x", relay.TensorType(dshape, "float32")) ndim = len(dshape) begin = begin if begin else [0] * ndim end = end if end else list(dshape) # target numpy result x_data = np.random.uniform(size=dshape).astype("float32") ref_res = tvm.topi.testing.strided_slice_python( x_data, begin, end, strides, slice_mode, axes=axes, ) if strides: z = relay.strided_slice(x, begin=begin, end=end, strides=strides, axes=axes, slice_mode=slice_mode) else: z = relay.strided_slice(x, begin=begin, end=end, axes=axes, slice_mode=slice_mode) func = relay.Function([x], z) func = run_infer_type(func) text = func.astext() assert "begin=" in text assert "end=" in text if output: assert func.body.checked_type == relay.ty.TensorType( output, "float32") if not test_ref: return for target, dev in tvm.testing.enabled_targets(): op_res = relay.create_executor( "graph", device=dev, target=target).evaluate(func)(x_data) tvm.testing.assert_allclose(op_res.numpy(), ref_res)
def expected(x, w1, w2, scale1, scale2, channels1, channels2): args = [x, w1, w2, scale1, scale2] w = relay.concatenate((w1, w2), axis=0) y = relay.nn.conv2d(x, w, channels=channels1 + channels2) y1 = relay.strided_slice(y, [0, 0], [None, channels1]) y2 = relay.strided_slice(y, [0, channels1], [None, channels1 + channels2]) y1 = relay.multiply(y1, scale1) y2 = relay.multiply(y2, scale2) y = relay.Tuple((y1, y2)) return relay.Function(args, y)
def expected(x, w1, w2, scale1, scale2, channels1, channels2): args = [x, w1, w2, scale1, scale2] w = relay.concatenate((w1, w2), axis=0) y = relay.nn.conv2d(x, w, channels=channels1 + channels2) y1 = relay.strided_slice(y, [0, 0], [None, channels1]) y2 = relay.strided_slice(y, [0, channels1], [None, channels1 + channels2]) y1 = relay.multiply(y1, scale1) y2 = relay.multiply(y2, scale2) y = relay.Tuple((y1, y2)) return relay.Function(args, y)
def expected(x, w, channels, repeat): args = [x, w] y = x for i in range(repeat): w_concat = relay.concatenate((w, w), axis=0) y = relay.nn.conv2d(y, w_concat, channels=channels*2) y1 = relay.strided_slice(y, [0, 0], [None, channels]) y2 = relay.strided_slice(y, [0, channels], [None, channels * 2]) y = relay.concatenate((y1, y2), axis=1) return relay.Function(args, y)
def expected(x, w, channels, repeat): args = [x, w] y = x for i in range(repeat): w_concat = relay.concatenate((w, w), axis=0) y = relay.nn.conv2d(y, w_concat, channels=channels * 2) y1 = relay.strided_slice(y, [0, 0], [None, channels]) y2 = relay.strided_slice(y, [0, channels], [None, channels * 2]) y = relay.concatenate((y1, y2), axis=1) return relay.Function(args, y)
def expected(dshape): x = relay.var("x", shape=dshape) slice1 = relay.strided_slice(x, begin=[0, 0], end=[dshape[1]//2, dshape[1]], strides=[1,1]) slice2 = relay.strided_slice(x, begin=[dshape[1]//2, 0], end=[dshape[0], dshape[1]], strides=[1,1]) out = relay.Tuple((slice1, slice2)) f0 = relay.Function([x], out) x = relay.var("x", shape=dshape) y = relay.Call(f0, [x]) return relay.Function([x], y)
def verify( dshape, begin_val, end_val, strides_val, output, slice_mode="end", test_ref=True, dtype="int32", ): x = relay.var("x", relay.TensorType(dshape, "float32")) ndim = len(dshape) begin_val = begin_val if begin_val else [0] * ndim end_val = end_val if end_val else list(dshape) if strides_val: if len(strides_val) == 1: strides_val = strides_val * ndim else: strides_val = [1] * ndim # target numpy result x_data = np.random.uniform(size=dshape).astype("float32") ref_res = tvm.topi.testing.strided_slice_python( x_data, begin_val, end_val, strides_val, slice_mode) data = [x_data, np.array(begin_val), np.array(end_val)] begin = relay.var("begin", relay.TensorType((len(begin_val), ), dtype)) end = relay.var("end", relay.TensorType((len(end_val), ), dtype)) func_params = [x, begin, end] if strides_val: data.append(np.array(strides_val)) strides = relay.var("strides", relay.TensorType((len(strides_val), ), dtype)) z = relay.strided_slice(x, begin=begin, end=end, strides=strides, slice_mode=slice_mode) func_params.append(strides) else: z = relay.strided_slice(x, begin=begin, end=end, slice_mode=slice_mode) func = relay.Function(func_params, z) params = {"begin": begin_val, "end": end_val, "strides": strides_val} func = run_infer_type(func) func2 = run_opt_pass( run_opt_pass(func, transform.DynamicToStatic(), params), transform.InferType()) assert isinstance(func2.body, relay.Call) assert func2.body.op == relay.op.get("strided_slice") verify_func(func2, [x_data], ref_res)
def verify(dshape, begin, end, strides, output, slice_mode="end", attr_const=True, test_ref=True, dtype="int32"): x = relay.var("x", relay.TensorType(dshape, "float32")) ndim = len(dshape) begin = begin if begin else [0] * ndim end = end if end else list(dshape) # target numpy result x_data = np.random.uniform(size=dshape).astype("float32") ref_res = tvm.topi.testing.strided_slice_python( x_data, begin, end, strides, slice_mode) if attr_const: begin = relay.const(begin, dtype=dtype) end = relay.const(end, dtype=dtype) if strides: if attr_const: strides = relay.const(strides, dtype=dtype) z = relay.strided_slice(x, begin=begin, end=end, strides=strides, slice_mode=slice_mode) else: z = relay.strided_slice(x, begin=begin, end=end, slice_mode=slice_mode) func = relay.Function([x], z) func = run_infer_type(func) text = func.astext() assert "begin=" in text assert "end=" in text if output: assert func.body.checked_type == relay.ty.TensorType( output, "float32") if not test_ref: return for target, ctx in ctx_list(): intrp = relay.create_executor("graph", ctx=ctx, target=target) op_res = intrp.evaluate(func)(x_data) tvm.testing.assert_allclose(op_res.asnumpy(), ref_res)
def before(dshape): x = relay.var("x", shape=dshape) slice1 = relay.strided_slice(x, begin=[0, 0], end=[dshape[1] // 2, dshape[1]], strides=[1, 1]) slice2 = relay.strided_slice(x, begin=[dshape[1] // 2, 0], end=[dshape[0], dshape[1]], strides=[1, 1]) out = relay.Tuple((slice1, slice2)) return relay.Function([x], out)
def expected(x, w1, w2, w3, w4, channels1, channels2, channels3, channels4): # use a fixed order of args so alpha equal check can pass args = [x, w1, w2, w3, w4] w = relay.concatenate((w1, w2, w4), axis=0) y = relay.nn.conv2d(x, w, channels=channels1 + channels2 + channels4) y1 = relay.strided_slice(y, [0, 0], [None, channels1]) y2 = relay.strided_slice(y, [0, channels1], [None, channels1 + channels2]) y3 = relay.nn.conv2d(x, w3) y4 = relay.strided_slice(y, [0, channels1 + channels2], [None, channels1 + channels2 + channels4]) y5 = relay.nn.max_pool2d(x) y = relay.Tuple((y1, y2, y3, y4, y5)) return relay.Function(args, y)
def expected(x, w1, w2, w3, w4, channels1, channels2, channels3, channels4): # use a fixed order of args so alpha equal check can pass args = [x, w1, w2, w3, w4] w = relay.concatenate((w1, w2, w4), axis=0) y = relay.nn.conv2d(x, w, channels=channels1 + channels2 + channels4) y1 = relay.strided_slice(y, [0, 0], [None, channels1]) y2 = relay.strided_slice(y, [0, channels1], [None, channels1 + channels2]) y3 = relay.nn.conv2d(x, w3) y4 = relay.strided_slice(y, [0, channels1 + channels2], [None, channels1 + channels2 + channels4]) y5 = relay.nn.max_pool2d(x) y = relay.Tuple((y1, y2, y3, y4, y5)) return relay.Function(args, y)
def expected(x, w1, w2, scale1, scale2, channels1, channels2): args = [x, w1, w2, scale1, scale2] w = relay.concatenate((w1, w2), axis=0) y = relay.nn.conv2d(x, w, channels=channels1 + channels2) y1 = relay.strided_slice( y, begin=[0, 0], end=[-1, channels1], strides=[1, 1], slice_mode="size" ) y2 = relay.strided_slice( y, begin=[0, channels1], end=[-1, channels2], strides=[1, 1], slice_mode="size" ) y1 = relay.multiply(y1, scale1) y2 = relay.multiply(y2, scale2) y = relay.Tuple((y1, y2)) return relay.Function(args, y)
def expected(x, w, channels, repeat): args = [x, w] y = x for i in range(repeat): w_concat = relay.concatenate((w, w), axis=0) y = relay.nn.conv2d(y, w_concat, channels=channels * 2) y1 = relay.strided_slice( y, begin=[0, 0], end=[-1, channels], strides=[1, 1], slice_mode="size" ) y2 = relay.strided_slice( y, begin=[0, channels], end=[-1, channels], strides=[1, 1], slice_mode="size" ) y = relay.concatenate((y1, y2), axis=1) return relay.Function(args, y)
def verify_any_strided_slice(data_shape, begin_shape, end_shape, strides_shape, data_np_shape, slice_mode="end", const_attrs=False): # Generate random numpy input data np_data = np.random.uniform(size=data_np_shape).astype('float32') np_begin = np.random.randint(2, size=begin_shape, dtype="int32") np_end = np.random.randint(5, 10, size=end_shape, dtype="int32") np_strides = np.random.randint(1, 2 if slice_mode == "size" else 3, size=strides_shape, dtype="int32") # target numpy result ref_res = topi.testing.strided_slice_python(np_data, np_begin, np_end, np_strides, slice_mode) # Relay Module mod = tvm.IRModule() data = relay.var('data', shape=data_shape, dtype='float32') if const_attrs: data = relay.var('data', shape=data_np_shape, dtype='float32') begin = relay.const(np_begin) end = relay.const(np_end) strides = relay.const(np_strides) args = [data] np_inputs = [np_data] else: begin = relay.var('begin', shape=begin_shape, dtype="int32") end = relay.var('end', shape=end_shape, dtype="int32") strides = relay.var('strides', shape=strides_shape, dtype="int32") args = [data, begin, end, strides] np_inputs = [np_data, np_begin, np_end, np_strides] y = relay.strided_slice(data, begin=begin, end=end, strides=strides, slice_mode=slice_mode) mod["main"] = relay.Function(args, y) for kind in ["debug", "vm"]: ex = relay.create_executor(kind, mod=mod, ctx=tvm.cpu(), target="llvm") result = ex.evaluate()(*np_inputs) tvm.testing.assert_allclose(result.asnumpy(), ref_res)
def get_graph(): x = relay.var("x", shape=(1, 2, 2, 4), dtype="int8") x = relay.strided_slice(x, begin=(0, 0, 0, 0), end=(1, 2, 2, 2)) x = infra.make_ethosu_identity(x) x = relay.reshape(x, newshape=(1, 2, 4)) x = infra.make_ethosu_identity(x) return relay.Function(relay.analysis.free_vars(x), x)
def callback(self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map) -> tvm.relay.Expr: split_input = post.args[0] split_begins = list() split_ends = list() section_begins_in_split_axis = self.get_section_begin_coords(post) for split_cord in section_begins_in_split_axis: # first begin is [0, 0, ... , 0] begin_shape = [ 0 for i in range(len(split_input.checked_type.shape)) ] begin_shape[post.attrs.axis] = split_cord split_begins.append(begin_shape) end_shape = list(split_input.checked_type.shape) # Only the split axis coordinate changes end_shape[post.attrs.axis] = split_cord split_ends.append(end_shape) # Coordinates needs to be shifted left because beginning # of the next section is the end of the previous split_ends = split_ends[1:] # Last section end is the shape of the tensor itself. split_ends.append(list(split_input.checked_type.shape)) strided_slices = list() for sb, se in zip(split_begins, split_ends): strided_slices.append(relay.strided_slice(split_input, sb, se)) return relay.Tuple(strided_slices)
def before(): x = relay.var("x", shape=(1, 32, 28, 28)) weight = relay.var('weight', shape=(32, 32, 3, 3)) y = relay.nn.conv2d(x, weight, channels=32, kernel_size=(3, 3), padding=(1, 1)) y = relay.strided_slice(y, begin=[0, 16], end=[None, None]) y = relay.Function(free_vars(y), y) return y
def conv2relay(node: Conv, batch_size): node2var = {} params = {} for term in node.inputs: for value in term: x = relay.var(name_hint=value.node.name, shape=(batch_size, *value.node.output_shape)) node2var[value.node] = x term_vars = [] for term in node.inputs: value_vars = [] for value in term: if value.begin == 0 and value.end == value.node.output_shape[0]: var = node2var[value.node] else: var = relay.strided_slice(node2var[value.node], begin=[0, value.begin, 0, 0], end=[batch_size, value.end, *value.node.output_shape[1:]]) value_vars.append(var) term_var = value_vars[0] for value_var in value_vars[1:]: term_var = term_var + value_var term_vars.append(term_var) if len(term_vars) > 1: x = relay.concatenate(term_vars, axis=1) else: x = term_vars[0] x = conv2d(x, node, params) node2var[node] = x fn = relay.Function(relay.analysis.free_vars(x), x) if tvm_minor_version() <= 6: return relay.Module.from_expr(fn), params else: return tvm.ir.IRModule.from_expr(fn), params
def before(): x = relay.var("x", shape=(1, 32, 28, 28)) weight = relay.var('weight', shape=(32, 32, 3, 3)) y = relay.nn.conv2d(x, weight, channels=32, kernel_size=(3, 3), padding=(1, 1)) y = relay.strided_slice(y, begin=[0, 16], end=[None, None]) y = relay.Function(analysis.free_vars(y), y) return y
def get_graph(): in_1 = relay.var("x", shape=(1, 16, 16, 8), dtype="int8") conv_1 = infra.make_ethosu_conv2d( ifm=in_1, ifm_channels=8, ofm_channels=8, kernel_shape=(1, 1), padding=(0, 0), strides=(1, 1), dilation=(1, 1), ifm_layout="NHWC", ofm_layout="NHWC", ) reshape = relay.reshape(conv_1, (1, 16, 16, 8)) strided_slice = relay.strided_slice(reshape, (0, 0, 0, 0), (1, 16, 16, 8)) conv_2 = infra.make_ethosu_conv2d( ifm=strided_slice, ifm_channels=8, ofm_channels=8, kernel_shape=(1, 1), padding=(0, 0), strides=(1, 1), dilation=(1, 1), ifm_layout="NHWC", ofm_layout="NHWC", ) return relay.Function(relay.analysis.free_vars(conv_2), conv_2)
def expected(): if not do_pad: return before() x = relay.var("x", shape=data_shape, dtype=dtype) if db or di: x_pad = relay.nn.pad(x, pad_width=((0, db), (0, 0), (0, 0), (0, di))) else: x_pad = x weight = relay.var("weight", shape=(kernel_shape), dtype=dtype) if di or do: weight_pad = relay.nn.pad(weight, pad_width=((0, 0), (0, 0), (0, di), (0, do))) else: weight_pad = weight y_pad = relay.nn.conv2d( x_pad, weight=weight_pad, channels=out_channel + do, kernel_size=(3, 3), padding=(1, 1), data_layout="NHWC", kernel_layout="HWIO", ) if db or do: y = relay.strided_slice(y_pad, begin=[0, 0, 0, 0], end=out_shape) else: y = y_pad y = relay.Function([x, weight], y) return y
def _pad_conv2d_NHWC(db, di, do, data, kernel, out_channel, new_attrs, output_tensor): # Pad batch size if db != 0: data = relay.nn.pad(data, pad_width=((0, db), (0, 0), (0, 0), (0, 0))) # Pad input channel if di != 0: data = relay.nn.pad(data, pad_width=((0, 0), (0, 0), (0, 0), (0, di))) kernel = relay.nn.pad(kernel, pad_width=((0, 0), (0, 0), (0, di), (0, 0))) # Pad output channel if do != 0: kernel = relay.nn.pad(kernel, pad_width=((0, 0), (0, 0), (0, 0), (0, do))) if do != 0: new_out_channel = out_channel + do new_attrs["channels"] = new_out_channel out = relay.nn.conv2d(data, kernel, **new_attrs) if db != 0 or do != 0: original_out_shape = [x.value for x in output_tensor.shape] out = relay.strided_slice(out, begin=[0, 0, 0, 0], end=original_out_shape) return out
def _get_func(): ifm = relay.var("ifm", shape=(1, 12, 12, 8), dtype="int8") conv1 = make_ethosu_conv2d( ifm=ifm, ifm_channels=8, ofm_channels=64, kernel_shape=(1, 1), padding=(0, 0), strides=(1, 1), dilation=(1, 1), activation="NONE", ifm_layout="NHWC", ofm_layout="NHWC", ) strided_slice = relay.strided_slice(conv1, [0, 0, 0, 0], [1, 6, 6, 128]) conv2 = make_ethosu_conv2d( ifm=strided_slice, ifm_channels=64, ofm_channels=16, kernel_shape=(3, 3), padding=(1, 1), strides=(1, 1), dilation=(1, 1), activation="NONE", ifm_layout="NHWC", ofm_layout="NHCWB16", ) func = relay.Function(relay.analysis.free_vars(conv2), conv2) func = run_opt_pass(func, relay.transform.InferType()) return func
def verify_strided_slice(dshape, begin, end, strides, mode): x = relay.var("x", relay.TensorType(dshape, "float32")) if mode == "size": strides = None z = relay.strided_slice(x, begin=begin, end=end, strides=strides, slice_mode=mode) func = relay.Function([x], z) x_data = np.random.uniform(size=dshape).astype("float32") verify_results(func, [x_data], "test_strided_slice", rtol=1e-5, atol=1e-5)
def before(): shape = (tvm.tir.const(10, "int64"), tvm.tir.const(1, "int64")) x = relay.var("x", shape=shape) concat = relay.concatenate([x,x], axis=-1) out = relay.strided_slice(concat, begin=[np.int64(0)], end=[np.int64(3)]) t = relay.Function(relay.analysis.free_vars(out), out) return relay.Function(relay.analysis.free_vars(out), out)
def get_graph(x_shape, begin, end, strides=None): x = relay.var("x", shape=(x_shape), dtype="float32") if strides: out = relay.strided_slice( x, relay.expr.const(begin, dtype="int32"), relay.expr.const(end, dtype="int32"), relay.expr.const(strides, dtype="int32"), ) else: out = relay.strided_slice( x, relay.expr.const(begin, dtype="int32"), relay.expr.const(end, dtype="int32"), ) f = relay.Function([x], out) return f, {"x": x_shape}, []
def get_graph(get_expected=False): x = relay.var("x", shape=(1, 2, 2, 4), dtype="int8") x = infra.make_ethosu_pooling(x, "MAX", (1, 1), 4, (1, 1), (0, 0)) x = relay.strided_slice(x, begin=[0, 0, 0, 0], end=[1, 2, 2, 2]) if not get_expected: x = infra.make_ethosu_identity(x) x = infra.make_ethosu_pooling(x, "MAX", (1, 1), 2, (1, 1), (0, 0)) return relay.Function(relay.analysis.free_vars(x), x)
def expected(): x = relay.var("x", shape=(1, 32, 28, 28)) weight = relay.var("weight") x = relay.layout_transform(x, "NCHW", "NCHW4c") y = relay.nn.conv2d(x, weight, channels=32, kernel_size=(3, 3), padding=(1, 1), data_layout="NCHW4c") y = relay.strided_slice(y, begin=[0, 4], end=[None, 8]) y = relay.layout_transform(y, "NCHW4c", "NCHW") y = relay.Function(free_vars(y), y) return y
def verify(dshape, begin, end, strides, output, test_ref=True): x = relay.var("x", relay.TensorType(dshape, "float32")) z = relay.strided_slice(x, begin=begin, end=end, strides=strides) func = relay.Function([x], z) func = relay.ir_pass.infer_type(func) text = func.astext() assert "begin=" in text assert "end=" in text if output: assert func.body.checked_type == relay.ty.TensorType(output, "float32") if not test_ref: return x_data = np.random.uniform(size=dshape).astype("float32") ref_res = topi.testing.strided_slice_python( x_data, begin, end, strides) for target, ctx in ctx_list(): intrp = relay.create_executor("graph", ctx=ctx, target=target) op_res = intrp.evaluate(func)(x_data) tvm.testing.assert_allclose(op_res.asnumpy(), ref_res)