def _verify(indices_shape, depth, on_value, off_value, axis, dtype): indices = relay.var("indices", relay.TensorType(indices_shape, "int32")) depth_var = relay.var("depth", relay.TensorType((), "int32")) on_value_const = relay.const(on_value) off_value_const = relay.const(off_value) out = relay.one_hot(indices, on_value_const, off_value_const, depth_var, axis, dtype) func = relay.Function([indices, depth_var], out) indices_np = np.random.randint(0, depth, size=indices_shape).astype("int32") out_np = tvm.topi.testing.one_hot(indices_np, on_value, off_value, depth, axis, dtype) for target, ctx in tvm.testing.enabled_targets(): if (target != 'cuda' ): #skip cuda because we don't have dynamic support for GPU for kind in ["vm", "debug"]: mod = tvm.ir.IRModule.from_expr(func) intrp = relay.create_executor(kind, mod=mod, ctx=ctx, target=target) out_relay = intrp.evaluate()( indices_np, np.array(depth).astype("int32")) tvm.testing.assert_allclose(out_relay.asnumpy(), out_np)
def _verify(indices_shape, depth, on_value, off_value, axis, dtype): indices = relay.var("indices", relay.TensorType(indices_shape, "int32")) depth_var = relay.const(depth) on_value_var = relay.var("on_value", relay.TensorType((), "int32")) off_value_var = relay.var("off_value", relay.TensorType((), "int32")) out = relay.one_hot(indices, on_value_var, off_value_var, depth_var, axis, dtype) params = { "on_value": on_value, "off_value": off_value, } func = relay.Function([indices, on_value_var, off_value_var], out) func2 = run_opt_pass( run_opt_pass(func, transform.DynamicToStatic(), params), transform.InferType()) zz = func2.body assert isinstance(zz, relay.Call) assert zz.op == relay.op.get("one_hot") indices_np = np.random.randint(0, depth, size=indices_shape).astype("int32") out_np = tvm.topi.testing.one_hot(indices_np, on_value, off_value, depth, axis, dtype) verify_func(func2, [indices_np], out_np)
def test_one_hot_grad(executor_kind, target, dev, index_dtype, val_dtype): indices_shape = (3, 4) depth = 5 axis = -1 inputs = [ np.random.randint(depth, size=indices_shape, dtype=index_dtype), np.array(np.random.randn() * 1e-5).astype(val_dtype), np.array(np.random.randn() * 1e-5).astype(val_dtype), ] test_inputs = inputs[1:] indices = relay.var("indices", shape=indices_shape, dtype=index_dtype) on_val = relay.var("on_val", shape=tuple(), dtype=val_dtype) off_val = relay.var("off_val", shape=tuple(), dtype=val_dtype) y = relay.one_hot(indices, on_val, off_val, depth, axis, val_dtype) f = relay.Function([indices, on_val, off_val], y) check_grad( f, inputs=inputs, test_inputs=test_inputs, target_devices=[(target, dev)], executor_kind=executor_kind, )
def test_one_hot_grad(): indices_shape = (3, 4) depth = 5 axis = -1 for indices_dtype in ["int32", "int64"]: for val_dtype in ["float32", "float64"]: inputs = [ np.random.randint(depth, size=indices_shape, dtype=indices_dtype), np.array(np.random.randn() * 1e-5).astype(val_dtype), np.array(np.random.randn() * 1e-5).astype(val_dtype), ] test_inputs = inputs[1:] indices = relay.var("indices", shape=indices_shape, dtype=indices_dtype) on_val = relay.var("on_val", shape=tuple(), dtype=val_dtype) off_val = relay.var("off_val", shape=tuple(), dtype=val_dtype) y = relay.one_hot(indices, on_val, off_val, depth, axis, val_dtype) f = relay.Function([indices, on_val, off_val], y) check_grad(f, inputs=inputs, test_inputs=test_inputs)
def _verify(indices_shape, depth, on_value, off_value, axis, dtype): indices = relay.var("indices", relay.TensorType(indices_shape, "int32")) on_value_const = relay.const(on_value) off_value_const = relay.const(off_value) out = relay.one_hot(indices, on_value_const, off_value_const, depth, axis, dtype) checked = run_infer_type(out) assert checked.checked_type == relay.ty.TensorType(_get_oshape(indices_shape, depth, axis), dtype) func = relay.Function([indices], out) indices_np = np.random.randint(0, depth, size=indices_shape).astype("int32") out_np = topi.testing.one_hot(indices_np, on_value, off_value, depth, axis, dtype) for target, ctx in ctx_list(): for kind in ["graph", "debug"]: intrp = relay.create_executor(kind, ctx=ctx, target=target) out_relay = intrp.evaluate(func)(indices_np) tvm.testing.assert_allclose(out_relay.asnumpy(), out_np)
def _verify(indices_shape, depth, on_value, off_value, axis, dtype): indices = relay.var("indices", relay.TensorType(indices_shape, "int32")) depth_var = relay.var("depth", relay.TensorType((), "int32")) on_value_const = relay.const(on_value) off_value_const = relay.const(off_value) out = relay.one_hot(indices, on_value_const, off_value_const, depth_var, axis, dtype) func = relay.Function([indices, depth_var], out) indices_np = np.random.randint(0, depth, size=indices_shape).astype("int32") out_np = tvm.topi.testing.one_hot(indices_np, on_value, off_value, depth, axis, dtype) for target, dev in tvm.testing.enabled_targets(): mod = tvm.ir.IRModule.from_expr(func) out_relay = relay.create_executor( executor_kind, mod=mod, device=dev, target=target).evaluate()(indices_np, np.array(depth).astype("int32")) tvm.testing.assert_allclose(out_relay.numpy(), out_np)
def _verify(indices_shape, depth, on_value, off_value, axis, dtype): indices = relay.var("indices", relay.TensorType(indices_shape, "int32")) on_value_const = relay.const(on_value) off_value_const = relay.const(off_value) out = relay.one_hot(indices, on_value_const, off_value_const, depth, axis, dtype) checked = run_infer_type(out) assert checked.checked_type == relay.ty.TensorType( _get_oshape(indices_shape, depth, axis), dtype) func = relay.Function([indices], out) indices_np = np.random.randint(0, depth, size=indices_shape).astype("int32") out_np = tvm.topi.testing.one_hot(indices_np, on_value, off_value, depth, axis, dtype) for target, dev in tvm.testing.enabled_targets(): out_relay = relay.create_executor( executor_kind, device=dev, target=target).evaluate(func)(indices_np) tvm.testing.assert_allclose(out_relay.numpy(), out_np)
def _execute(self): self.node_dict = {} # self.node_dict['1'] = relay.const(np.zeros((1, 128)), dtype='int32') gelu_a = relay.var('gelu_a', shape=()) gelu_b = relay.var('gelu_b', shape=()) gelu_c = relay.var('gelu_c', shape=()) gelu_d = relay.var('gelu_d', shape=()) gelu_e = relay.var('gelu_e', shape=()) self.node_dict['1'] = relay.var('input.1', shape=(1,128), dtype='int32') self.node_dict['2'] = relay.var('input.2', shape=(1,128), dtype='int32') for gnode in self.graph: name = gnode['name'] op_type = gnode['op_type'] attrs = gnode['attrs'] del attrs['A_shape'] del attrs['O_shape'] inputs = gnode['inputs'] if op_type == 'Const': arr = np.zeros(attrs['shape'], dtype=np.int32) y = relay.const(arr, dtype='int32') elif op_type == 'expand_dims': x = get_input(self.node_dict, self.params, inputs[0]) y = relay.expand_dims(x, attrs['axis'], attrs['num_newaxis']) elif op_type == 'reshape': x = get_input(self.node_dict, self.params, inputs[0]) y = relay.reshape(x, attrs['newshape']) elif op_type == 'take': data = get_input(self.node_dict, self.params, inputs[0]) indices = get_input(self.node_dict, self.params, inputs[1]) y = relay.take(data, indices, axis=attrs['axis'][0], mode=attrs['mode']) elif op_type == 'one_hot': x = get_input(self.node_dict, self.params, inputs[0]) cc1 = get_input(self.node_dict, self.params, inputs[1]) cc2 = get_input(self.node_dict, self.params, inputs[2]) y = relay.one_hot(x, cc1, cc2, **attrs) elif op_type == 'strided_slice': x = get_input(self.node_dict, self.params, inputs[0]) y = relay.strided_slice(x, **attrs) elif op_type == 'mean': x = get_input(self.node_dict, self.params, inputs[0]) y = relay.mean(x, axis=attrs['axis'], exclude=attrs['exclude'], keepdims=attrs['keepdims']) elif op_type == 'nn.dense': x = get_input(self.node_dict, self.params, inputs[0]) weight = get_input(self.node_dict, self.params, inputs[1]) y = relay.nn.dense(x, weight, units=attrs['units'][0]) elif op_type == 'add': x1 = get_input(self.node_dict, self.params, inputs[0]) x2 = get_input(self.node_dict, self.params, inputs[1]) y = relay.add(x1, x2) elif op_type == 'subtract': x1 = get_input(self.node_dict, self.params, inputs[0]) x2 = get_input(self.node_dict, self.params, inputs[1]) y = relay.subtract(x1, x2) elif op_type == 'multiply': x1 = get_input(self.node_dict, self.params, inputs[0]) x2 = get_input(self.node_dict, self.params, inputs[1]) y = relay.multiply(x1, x2) elif op_type == 'power': x1 = get_input(self.node_dict, self.params, inputs[0]) x2 = get_input(self.node_dict, self.params, inputs[1]) y = relay.power(x1, x2) elif op_type == 'transpose': x = get_input(self.node_dict, self.params, inputs[0]) y = relay.transpose(x, **attrs) elif op_type == 'tanh': x = get_input(self.node_dict, self.params, inputs[0]) y = relay.tanh(x) elif op_type == 'squeeze': x = get_input(self.node_dict, self.params, inputs[0]) y = relay.squeeze(x, **attrs) elif op_type == 'nn.batch_matmul': x1 = get_input(self.node_dict, self.params, inputs[0]) x2 = get_input(self.node_dict, self.params, inputs[1]) y = relay.nn.batch_matmul(x1, x2) elif op_type == 'nn.softmax': x = get_input(self.node_dict, self.params, inputs[0]) y = relay.nn.softmax(x, **attrs) elif op_type == 'gelu': x = get_input(self.node_dict, self.params, inputs[0]) y = x * gelu_a * (gelu_b + relay.tanh( ( gelu_c * (x + gelu_d * relay.power(x, gelu_e))))) else: import pdb; pdb.set_trace() print( 'not supported op %s ' % op_type) self.node_dict[name] = y output_name = self.output_node_ids[0] output = self.node_dict[output_name] inputs = relay.analysis.free_vars(output) # inputs = [self.node_dict['1'], self.node_dict['2']] func = relay.Function(inputs, output) mod = tvm.IRModule() mod['main'] = func with relay.build_config(opt_level=0): graph, lib, params = relay.build(mod, 'llvm', params={}) self.m = graph_runtime.create(graph, lib, tvm.cpu())