def run(zero_point, scale):
qparams = create_qparams(QuantMode.ASYMMERTIC, test_dtype, scale, zero_point)
inp_data = np.random.uniform(low=-512.0, high=512.0, size=(1, 32, 32, 32))
inp = tensor(inp_data, dtype=np.float32)
# test forward
oup = fake_quant_tensor(inp, qparams).numpy()
oup_gt = fake_quant_tensor_gt(inp, scale, zero_point, qmin, qmax).numpy()
assert np.allclose(oup, oup_gt)
assert oup.shape == oup_gt.shape
# test backward
x = tensor(inp_data, dtype=np.float32)
with Grad() as grad:
grad.wrt(x, callback=_save_to(x))
y = fake_quant_tensor(x, qparams)
grad(y, tensor(F.ones_like(x)))
x1 = tensor(inp_data, dtype=np.float32)
with Grad() as grad:
grad.wrt(x1, callback=_save_to(x1))
y1 = fake_quant_tensor_gt(x1, scale, zero_point, qmin, qmax)
grad(y1, tensor(F.ones_like(x1)))
assert np.allclose(x.grad.numpy(), x1.grad.numpy())
assert make_shape_tuple(x.grad.shape) == make_shape_tuple(x1.grad.shape)
# test nan
x = F.full((1, 32, 3, 3), np.nan)
y = fake_quant_tensor(x, qparams).numpy()
assert np.isnan(y).all()
def test_ShuffleRNG():
g = []
def cb(grad):
g.append(grad)
n, m = 6, 3
arr = np.arange(n * m)
out0 = Tensor(arr, dtype="float32")
with Grad() as grad:
grad.wrt(out0, callback=cb)
random.shuffle(out0)
grad(out0, F.ones_like(out0))
m1 = RNG(seed=111, device="xpu0")
m2 = RNG(seed=111, device="xpu1")
m3 = RNG(seed=222, device="xpu0")
out1 = Tensor(arr, dtype="float32", device="xpu0")
out2 = Tensor(arr, dtype="float32", device="xpu1")
out3 = Tensor(arr, dtype="float32", device="xpu0")
m1.shuffle(out1)
m2.shuffle(out2)
m3.shuffle(out3)
np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6)
assert out1.device == "xpu0" and out2.device == "xpu1"
assert not (out1.numpy() == out3.numpy()).all()
out = Tensor(arr, dtype="float32").reshape(n, m)
m1.shuffle(out)
out_shp = out.shape
if isinstance(out_shp, tuple):
assert out_shp == (n, m)
else:
assert all(out.shape.numpy() == np.array([n, m]))
def test_tqt():
g = []
def cb(grad):
g.append(grad)
x = np.random.randint(-128, 128, size=(1, 2, 3, 4)).astype("float32")
s = np.random.rand(1) - 1
g_y = np.ones(shape=(1, 2, 3, 4), dtype="float32")
n = TQT_numpy(-127, 127)
y_np = n.forward(x, s)
g_x_np, g_s_np = n.backward(g_y)
x = mge.tensor(x, dtype="float32")
s = mge.tensor(s, dtype="float32")
g_y = mge.tensor(g_y, dtype="float32")
with Grad() as grad:
grad.wrt(x, s, callback=cb)
y = tqt_forward(-127, 127, x, s)
grad(y, g_y)
g_x, g_s = g
np.testing.assert_allclose(y.numpy(), y_np, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(g_x.numpy(), g_x_np, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(g_s.numpy(), g_s_np, rtol=5e-5, atol=5e-5)
def worker():
rank = dist.get_rank()
if rank == 0:
with Grad() as grad:
x = as_tensor(x_np)
grad.wrt(x, callback=save_to(x))
# need a placeholder to trace operator
remote_send(x, 1)
recv_x = remote_recv(1)
y = recv_x * recv_x
grad([y], [as_tensor(np.ones_like(x_np))])
np.testing.assert_almost_equal(x.grad.numpy(), x.numpy() * 2)
elif rank == 1:
with Grad() as grad:
recv_x = remote_recv(0)
remote_send(recv_x, 0)
grad([], [])
def test_resize():
x_np = np.random.rand(3, 3, 32, 32).astype("float32")
x = mge.Tensor(x_np)
grad = Grad().wrt(x, callback=save_to(x))
y = F.resize(x, (16, 16))
grad(y, F.ones_like(y))
np.testing.assert_equal(np.ones(x_np.shape, dtype=np.float32) / 4, x.grad.numpy())
def test_2nd_grad():
x_np = np.random.rand(10).astype("float32")
x = as_tensor(x_np)
ones = as_tensor(np.ones_like(x_np))
with Grad("grad2") as grad2:
with Grad("grad") as grad:
grad2.wrt(x, callback=save_to(x))
grad.wrt(x, callback=save_to(x))
y = cos(x)
grad(y, ones)
z = x.grad
np.testing.assert_almost_equal(x.grad.numpy(), -np.sin(x_np), decimal=5)
x.grad = None
grad2(z, ones)
np.testing.assert_almost_equal(x.grad.numpy(), -np.cos(x_np), decimal=5)
def test_AxisAddRemove():
x_np = np.random.rand(1, 5).astype("float32")
x = TensorWrapper(x_np)
grad = Grad().wrt(x, callback=save_to(x))
y = F.squeeze(F.expand_dims(x, 2), 0)
grad(y, F.ones_like(y))
np.testing.assert_equal(np.array([[1, 1, 1, 1, 1]], dtype=np.float32),
x.grad.numpy())
def test_grad():
x_np = np.random.rand(10).astype("float32")
x = as_tensor(x_np)
with Grad() as grad:
grad.wrt(x, callback=save_to(x))
y = cos(x)
grad(y, as_tensor(np.ones_like(x_np)))
np.testing.assert_almost_equal(x.grad.numpy(), -np.sin(x_np))
def test_interpolate_fastpath():
x_np = np.random.rand(3, 3, 32, 32).astype("float32")
x = mge.Tensor(x_np)
with Grad() as grad:
grad.wrt(x, callback=save_to(x))
y = F.vision.interpolate(x, size=(16, 16), mode="bilinear")
grad(y, F.ones_like(y))
np.testing.assert_equal(np.ones(x_np.shape, dtype=np.float32) / 4, x.grad.numpy())
def test_grad_2():
x_np = np.random.rand(10).astype("float32")
x = as_tensor(x_np)
with Grad() as grad:
grad.wrt(x, callback=save_to(x))
y = mul(x, x)
y = mul(y, y)
grad(y, as_tensor(np.ones_like(x_np)))
np.testing.assert_almost_equal(x.grad.numpy(), 4 * x_np ** 3, decimal=6)
def test_lsq():
g = []
def cb(grad):
g.append(grad)
# FIXME: use random number when LSQ is fixed
# x = np.random.randint(-128, 128, size=(1, 2, 3, 4)).astype("float32")
# s = np.random.rand(1)
x = np.array(
[
[
[
[4.0, 38.0, -121.0, 38.0],
[15.0, -115.0, -112.0, 24.0],
[23.0, -65.0, 109.0, -115.0],
],
[
[-66.0, -90.0, -45.0, -101.0],
[68.0, -98.0, 108.0, -79.0],
[54.0, 63.0, -10.0, -50.0],
],
]
],
dtype="float32",
)
s = np.array([0.02918224], dtype="float32")
eps = np.array([1e-5], dtype="float32")
s = np.abs(s) if np.abs(s) > eps else eps
zero_point = np.array([1.0], dtype="float32")
grad_s = np.array([2.0], dtype="float32")
g_y = np.ones(shape=(1, 2, 3, 4), dtype="float32")
n = LSQ_numpy(-127, 127)
y_np = n.forward(x, s, zero_point, grad_s)
g_x_np, g_s_np = n.backward(g_y)
x = mge.tensor(x, dtype="float32")
s = mge.tensor(s, dtype="float32")
zero_point = mge.tensor(zero_point, dtype="float32")
grad_s = mge.tensor(grad_s, dtype="float32")
g_y = mge.tensor(g_y, dtype="float32")
with Grad() as grad:
grad.wrt(x, s, callback=cb)
y = lsq_forward(-127, 127, x, s, zero_point, grad_s)
grad(y, g_y)
g_x, g_s = g
np.testing.assert_allclose(y.numpy(), y_np, rtol=1e-7, atol=1e-7)
np.testing.assert_allclose(g_x.numpy(), g_x_np, rtol=1e-7, atol=1e-7)
np.testing.assert_allclose(g_s.numpy(), g_s_np, rtol=5e-7, atol=5e-7)
def test_identity():
x_np = np.random.rand(10).astype("float32")
x = mge.Tensor(x_np)
dy_np = np.random.rand(*x.shape).astype("float32")
dy = mge.Tensor(dy_np)
grad = Grad().wrt(x, callback=save_to(x))
(y,) = apply(Identity(), x)
grad(y, dy)
np.testing.assert_array_equal(x.grad.numpy(), dy_np)
def test_dot():
x = np.random.rand(2, 2).astype("float32")
x = mge.Tensor(x)
u = F.ones((2,))
v = F.ones((2,))
with Grad() as grad:
grad.wrt(x, callback=save_to(x))
def f(x):
return F.dot(u, F.matmul(x, v))
y = f(x)
grad(y, F.ones_like(y))
np.testing.assert_equal(np.ones((2, 2), dtype=np.float32), x.grad.numpy())