def lstm_driver(): bg = tvm.placeholder((10, 1), name='bg') Ug = tvm.placeholder((10, 10), name='Ug') Vg = tvm.placeholder((10, 10), name='Vg') bi = tvm.placeholder((10, 1), name='bi') Ui = tvm.placeholder((10, 10), name='Ui') Vi = tvm.placeholder((10, 10), name='Vi') bf = tvm.placeholder((10, 1), name='bf') Uf = tvm.placeholder((10, 10), name='Uf') Vf = tvm.placeholder((10, 10), name='Vf') bo = tvm.placeholder((10, 1), name='bo') Uo = tvm.placeholder((10, 10), name='Uo') Vo = tvm.placeholder((10, 10), name='Vo') x = tvm.placeholder((10, 1), name='x') s = tvm.placeholder((10, 1), name='s') h = tvm.placeholder((10, 1), name='h') s2, h2 = lstm_cell(bg, Ug, Vg, bi, Ui, Vi, bf, Uf, Vf, bo, Uo, Vo, x, s, h) print(s2, h2) sout = tvm.create_schedule([s2.op, h2.op]) mout = tvm.build( sout, [s2, h2, bg, Ug, Vg, bi, Ui, Vi, bf, Uf, Vf, bo, Uo, Vo, x, s, h], 'llvm') print(mout) ones = topi.full_like(s2, 1.0) jac = tvm.ir_pass.JacobianRecursive(s2, x, ones) print(jac)
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 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 example2_build(): npoints = 20 x = tvm.placeholder((1, )) k = tvm.placeholder((2, )) y = tvm.compute((1, ), lambda i: x[i] * x[i] * k[0] + x[i] * k[1], name='y') ones = topi.full_like(y, 1.0) [dy] = list(tvm.ir_pass.JacobianRecursive(y, [x], ones)) # print('jac',type(jac),jac[0]) sdy = tvm.create_schedule(dy.op) sy = tvm.create_schedule(y.op) my = tvm.build(sy, [y, x, k]) mdy = tvm.build(sdy, [dy, x, k]) xs = np.linspace(-10.0, 10.0, 20) ys = [] dys = [] for xi in xs: k_in = tvm.nd.array(np.array([4.0, 2.0]).astype(k.dtype)) y_out = tvm.nd.empty(get_shape(y), y.dtype) my(y_out, tvm.nd.array(np.array([xi]).astype(x.dtype)), k_in) ys.append(y_out.asnumpy()) dy_out = tvm.nd.empty(get_shape(y), dy.dtype) mdy(dy_out, tvm.nd.array(np.array([xi]).astype(x.dtype)), k_in) dys.append(dy_out.asnumpy()) # print(xi, y_out.asnumpy(), dy_out.asnumpy()) print(xs, ys, dys)
def example1_build(): x = tvm.placeholder((1, )) k = tvm.placeholder((1, )) y = tvm.compute((1, ), lambda i: x[i] * x[i] * k[0], name='y') [dy] = tvm.ir_pass.JacobianRecursive(y, [x], topi.full_like(y, 1.0)) sdy = tvm.create_schedule(dy.op) mdy = tvm.build(sdy, [dy, x, k]) dy_out = tvm.nd.empty(get_shape(y), tvm.float32) mdy(dy_out, tvm.nd.array(np.array([1.0]).astype(x.dtype)), tvm.nd.array(np.array([4.0]).astype(k.dtype))) print(dy_out)
def check_device(device, host="llvm"): ctx = tvm.context(device, 0) if not tvm.runtime.enabled(host): return if not ctx.exist: print("skip because %s is not enabled.." % device) return sout = te.create_schedule(out.op) mout = tvm.build(sout, [out] + inputs) out_shape = get_const_tuple(out.shape) l, h = data_range input_data = [ tvm.nd.array( np.random.uniform(l, h, size=get_const_tuple( input.shape)).astype(input.dtype)) for input in inputs ] ones = topi.full_like(out, 1.0) # we provide head to sum and reduce the output dimension, # which equals to grad(out.sum(), inputs) grads = te.gradient(out, inputs, head=ones) grad_sched = te.create_schedule([grad.op for grad in grads]) mgrad = tvm.build(grad_sched, list(grads) + inputs) # print(tvm.lower(grad_sched, list(grads) + inputs, simple_mode=True)) grad_data = [ tvm.nd.empty(get_const_tuple(i.shape), g.dtype) for i, g in zip(inputs, grads) ] mgrad(*grad_data, *input_data) g_res = [g.asnumpy() for g in grad_data] if desired_grads: assert isinstance(desired_grads, list) for actual, desired in zip(g_res, desired_grads): assert_allclose(actual, desired, rtol=0.1, atol=1e-2) else: def forward(*in_data): out_data = tvm.nd.empty(out_shape, out.dtype) mout(out_data, *[tvm.nd.array(d) for d in list(in_data)]) return out_data.asnumpy().sum() check_numerical_grads(forward, [d.asnumpy() for d in input_data], g_res)
def ones_like(attrs, inputs, output_type, target): assert len(inputs) == 1 return [topi.full_like(inputs[0], 1.0)]
def zeros_like_compute(attrs, inputs, output_type, target): assert len(inputs) == 1 return [topi.full_like(inputs[0], 0.0)]
def demo_conv2d(): lrate = 0.1 nbatches = 100 # batches to train num_classes = 10 batch_size = 10 img_h = 28 img_w = 28 img_c = 1 f1_c = 4 f2_c = 5 f3_units = 16 x = tvm.placeholder((batch_size, img_h, img_w, img_c), name='x') y = tvm.placeholder((batch_size, num_classes), name='y') print('Block1') w1 = tvm.placeholder((3, 3, img_c, f1_c), name='w1') b1 = tvm.placeholder((f1_c, ), name='b1') t = topi.nn.conv2d(x, w1, 1, 0, layout='NHWC', out_dtype=tvm.float32) t = t + topi.broadcast_to(b1, (batch_size, 1, 1, f1_c)) print('Block1: after-biasing shape is', get_shape(t)) t = topi.nn.pool(t, [2, 2], [2, 2], [0, 0, 0, 0], 'max', layout='NHWC') print('Block1: after-pooling shape is', get_shape(t)) t = topi.nn.relu(t) print('Block1: after-relu shape is', get_shape(t)) print('Block2') w2 = tvm.placeholder((3, 3, f1_c, f2_c), name='w2') b2 = tvm.placeholder((f2_c, ), name='b2') t = topi.nn.conv2d(t, w2, 1, 0, layout='NHWC', out_dtype=tvm.float32) t = t + topi.broadcast_to(b2, (batch_size, 1, 1, f2_c)) print('Block2: after-biasing shape is', get_shape(t)) t = topi.nn.pool(t, [2, 2], [2, 2], [0, 0, 0, 0], 'max', layout='NHWC') print('Block2: after-pooling shape is', get_shape(t)) t = topi.nn.relu(t) print('Block2: after-relu shape is', get_shape(t)) t = topi.nn.flatten(t) print('Block2: after-flattern shape is', get_shape(t)) print('Block3') w3 = tvm.placeholder((f3_units, get_shape(t)[1])) b3 = tvm.placeholder((f3_units, )) t = topi.nn.dense(t, w3, b3) print('Block3: after-dense shape is', get_shape(t)) print('Block4') w4 = tvm.placeholder((num_classes, get_shape(t)[1])) b4 = tvm.placeholder((num_classes, )) t = topi.nn.dense(t, w4, b4) print('Block4: after-dense shape is', get_shape(t)) t = topi.nn.relu(t) p = topi.argmax(t, axis=1) # TODO: check the correctnesss of the log_softmax expression # TODO: figure out the difference between it and standard cross-entropy loss l = -topi.sum(y * topi.nn.log_softmax(t)) / batch_size print('Block4: loss shape is', get_shape(l)) ones = topi.full_like(l, 1.0) #[dl_dw1,dl_db1,dl_dw2,dl_db2,dl_dw3,dl_db3,dl_dw4,dl_db4] params = [w1, b1, w2, b2, w3, b3, w4, b4] dl = list(tvm.ir_pass.JacobianRecursive(l, params, ones)) assert len(params) == len(dl) print('dl_dw1 weight is', get_shape(params[0])) sdl = tvm.create_schedule([p.op for p in [x, y, l] + params + dl]) mdl = tvm.build(sdl, [x, y, l] + params + dl) print('Train+Inference module', mdl) # sl = tvm.create_schedule([l.op]) # ml = tvm.build(sdl, [x,y] + params + [l]) # print('Inference module',ml) state = {} for p in params: state.update({ p: tvm.nd.array( np.random.uniform(-1.0, 1.0, size=get_shape(p)).astype(np.float32)) }) grads = {} for p, g in zip(params, dl): grads.update({p: tvm.nd.empty(get_shape(g))}) for ib in range(nbatches): b = range(ib * batch_size, (ib + 1) * batch_size) tx = tvm.nd.array(mnist_img(b)) ty = tvm.nd.array(mnist_cls_oh(b)) tl = tvm.nd.empty(shape=(), dtype=tvm.float32) print('Entering') mdl(*([tx, ty, tl] + list(state.values()) + list(grads.values()))) print('Done', 'loss', tl.asnumpy()) state2 = {} for p in params: state2.update({ p: tvm.nd.array(state[p].asnumpy() - lrate * grads[p].asnumpy()) }) state = state2
def ones_like(attrs, inputs, output_type, target): assert len(inputs) == 1 return [topi.full_like(inputs[0], 1.0)]
def zeros_like_compute(attrs, inputs, output_type, target): assert len(inputs) == 1 return [topi.full_like(inputs[0], 0.0)]
import tvm import topi x = tvm.te.placeholder((32, 3, 28, 28), name='x') w1 = tvm.te.placeholder((10, 3, 3, 3), name='w1') w2 = tvm.te.placeholder((10, 10, 3, 3), name='w2') z1 = topi.nn.conv2d(x, w1, 1, 1, 1) z2 = topi.nn.conv2d(z1, w2, 1, 1, 1) y = topi.sum(z2) # produce gradients [dw1] = tvm.te.gradient(y, [w1]) print(type(dw1)) # produce Jacobians [jw1, jw2] = tvm.te.gradient(z2, [w1, w2]) # produce gradients, the head adjoint for z2 is provided manually [dw1, dw2] = tvm.te.gradient(z2, [w1, w2], topi.full_like(z2, 1.0))