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)
