def verify_full(shape, dtype, fill_value): A = tvm.placeholder(shape, dtype=dtype, name="A") B = topi.full_like(A, fill_value=fill_value) C = topi.full(shape=shape, dtype=dtype, fill_value=fill_value) s1 = tvm.create_schedule([B.op]) s2 = tvm.create_schedule([C.op]) @memoize("topi.tests.test_topi_full") def get_ref_data(): return np.full(shape, fill_value, dtype) np_nd = get_ref_data() def check_device(device): if not tvm.module.enabled(device): print("Skip because %s is not enabled" % device) return ctx = tvm.context(device, 0) out = tvm.nd.array(np.zeros(shape, dtype=dtype), ctx) f = tvm.build(s1, [A, B], device, name="full_like") f(tvm.nd.array(np.zeros(shape, dtype), ctx), out) np.testing.assert_allclose(out.asnumpy(), np_nd, rtol=1e-5) f = tvm.build(s2, [C], device, name="full") f(out) np.testing.assert_allclose(out.asnumpy(), np_nd, rtol=1e-5) for device in ["llvm"]: check_device(device)
def verify_full(shape, dtype, fill_value): A = tvm.placeholder(shape, dtype=dtype, name="A") B = topi.full_like(A, fill_value=fill_value) C = topi.full(shape=shape, dtype=dtype, fill_value=fill_value) s1 = tvm.create_schedule([B.op]) s2 = tvm.create_schedule([C.op]) @memoize("topi.tests.test_topi_full") def get_ref_data(): return np.full(shape, fill_value, dtype) np_nd = get_ref_data() def check_device(device): if not tvm.module.enabled(device): print("Skip because %s is not enabled" % device) return ctx = tvm.context(device, 0) out = tvm.nd.array(np.zeros(shape, dtype=dtype), ctx) f = tvm.build(s1, [A, B], device, name="full_like") f(tvm.nd.array(np.zeros(shape, dtype), ctx), out) tvm.testing.assert_allclose(out.asnumpy(), np_nd, rtol=1e-5) f = tvm.build(s2, [C], device, name="full") f(out) tvm.testing.assert_allclose(out.asnumpy(), np_nd, rtol=1e-5) for device in ["llvm"]: check_device(device)
def check_device(device): with tvm.target.create(device): ctx = tvm.context(device, 0) if not ctx.exist: print("Skip because %s is not enabled" % device) return print("Running on target: %s" % device) data = tvm.placeholder((2, 1, 2, 4), 'int8', 'data') w = tvm.placeholder((3, 1, 2, 2), 'int8', 'w') conv1 = topi.nn.conv2d(data, w, 1, 0, 1, out_dtype='int32') zeros = topi.full((2, 3, 1, 3), 'int32', tvm.const(0, dtype='int32')) gt = topi.greater_equal(conv1, zeros) one = topi.full((2, 3, 1, 3), 'int32', tvm.const(1, dtype='int32')) two = topi.full((2, 3, 1, 3), 'int32', tvm.const(2, dtype='int32')) where = topi.where(gt, one, two) add = topi.add(conv1, where) outs = [add] s = topi.generic.schedule_conv2d_nchw(outs) tvm.build(s, [data, w, add], target=backend)
def ones_compute(attrs, inputs, output_type, target): assert not inputs return [topi.full(output_type.shape, output_type.dtype, 1.0)]
def zeros_compute(attrs, inputs, output_type): assert not inputs return [topi.full(output_type.shape, output_type.dtype, 0.0)]
new_s = new_s * f + g * j new_h = topi.tanh(new_s) * o logits = topi.matmul(new_h, weights[8]) + weights[9] # compute accuracy pred = topi.nn.softmax(logits) correct_pred = topi.equal(topi.argmax(y, 1), topi.argmax(pred, 1)) accuracy = topi.sum(correct_pred.astype('float32')) / batch_size # Define loss and optimizer loss = topi.sum(-topi.sum(y * topi.nn.log_softmax(logits), axis=1)) / batch_size head = topi.full((1, ), 'float32', 1.0) gradients = list(tvm.differentiate(topi.reshape(loss, (1, )), weights, head)) new_weights = [w - lr * g for (w, g) in zip(weights, gradients)] # Define model sched = tvm.create_schedule([loss.op, accuracy.op] + [x.op for x in new_weights]) parallel_schedule(sched) train_model = tvm.build(sched, [x, y, s, h, loss, accuracy, *weights, *new_weights]) # Define variables for input to graph train_weights = [ tvm.ndarray.array(op[0](**{ op[1]: x }).astype('float32')) for (op, x) in zip(inits, sizes)
def ones_compute(attrs, inputs, output_type): assert len(inputs) == 1 return [topi.full(output_type.shape, output_type.dtype, 1.0)]