def test_backward(self): n, c, H, W = 2, 3, 5, 5 o, i, h, w = 4, c, 3, 3 a = np.random.rand(n, c, H, W).astype(np.float32) b = np.random.rand(o, i, h, w).astype(np.float32) da = np.random.rand(n, o, H, W).astype(np.float32) dx = jt.mkl_ops.mkl_conv_backward_x(b, da, H, W, 1, 1, 1, 1, 1, 1).data dw = jt.mkl_ops.mkl_conv_backward_w(a, da, h, w, 1, 1, 1, 1, 1, 1).data a_jt = jt.array(a) b_jt = jt.array(b) with jt.flag_scope(enable_tuner=0, # compile_options={"test_mkl_conv":1} ): c_jt = conv(a_jt, b_jt, 1, 1) * da gs = jt.grad(c_jt, [a_jt, b_jt]) gs.append(c_jt) jt.fetch_sync(gs) dx_jt = gs[0].data dw_jt = gs[1].data with jt.log_capture_scope( log_v=10, log_vprefix="tuner_manager=100,var_relay=100", enable_tuner=1, compile_options={"test_mkl_conv": 2}) as rawlogs: gs_tune = jt.grad(c_jt, [a_jt, b_jt]) jt.fetch_sync(gs_tune) dx_jt_tune = gs_tune[0].data dw_jt_tune = gs_tune[1].data logs = find_log_with_re( rawlogs, "Run tuner conv: confidence\\((20)\\) candidates\\((.*)\\)$") assert len(logs) == 2, len(logs) assert logs[0][0] == "20", "confidence of reorder should be 20" candidates = simple_parser(logs[0][1]) assert candidates == {"relay0": [1, 0]}, candidates logs = find_log_with_re(rawlogs, r"get_relay_src([\s\S]*)") assert len(logs) == 2 assert "@relay_op" in logs[0] assert "@relay_op" in logs[1] assert np.max(dx_jt - dx) < 1e-5 and np.max(dw_jt - dw) < 1e-5 assert np.max(dx_jt_tune - dx) < 1e-5 and np.max(dw_jt_tune - dw) < 1e-5
def test_pad(self): size = 10 lpad = 3 rpad = 4 a = jt.random([size]) b = a.reindex([size+lpad+rpad], [f"i0-{lpad}"], -1) na, nb = jt.fetch_sync([a, b]) assert (nb[lpad:lpad+size]==na).all() assert (nb[:lpad]==-1).all() assert (nb[-rpad:]==-1).all()
def test_matmul(self): size = 10 a = jt.random([size,size]) b = jt.random([size,size]) cc = a.reindex([size,size,size],["i0","i1"]) * \ b.reindex([size,size,size],["i1","i2"]) c = cc.sum(dim=1) na, nb, nc = jt.fetch_sync([a, b, c]) assert is_fused(cc) assert not is_fused(c) check_fused(len(a.shape)) npc = np.matmul(na,nb) assert np.allclose(npc, nc)
def test_conv_transpose_grad(self): N,H,W,C = 1,5,5,2 Kh, Kw, Kc = 3, 3, 2 x = jt.random([N,H,W,C]) w = jt.random([Kh,Kw,C,Kc]) y, yy = conv_transpose(x, w) mask = jt.random(y.shape) loss = (y*mask).sum() dx, dw = jt.grad(loss, [x, w]) jdx, jdw = jt.fetch_sync([dx, dw]) check_fused(len(x.shape)) nmask = mask.data _, (ndx, ndw) = ngrad(lambda args: \ (conv_transpose_naive(args[0], args[1])*nmask).sum(), [np.float64(x.data), np.float64(w.data)], 1e-7) assert np.allclose(ndx, jdx), (ndx, jdx, ndx-jdx) assert np.allclose(ndw, jdw), (ndw, jdw)
def test(self): a = jt.random([10]) b = jt.code(a.shape, a.dtype, [a], cpu_src=''' for (int i=0; i<in0_shape0; i++) @out(i) = @in0(i)*@in0(i)*2; ''', cpu_grad_src = [''' for (int i=0; i<in0_shape0; i++) { @out(i) = @dout(i)*@in0(i)*4; } ''']) na, nb = jt.fetch_sync([a,b]) assert np.allclose(na*na*2, nb) c = jt.random([10]) da = jt.grad(c*b, a) assert np.allclose(c.data*na*4, da.data), (c.data*na*4, da.data)