def test_backward(self):
np.random.seed(0)
jt.set_seed(3)
model = Model()
SGD = jt.nn.SGD(model.parameters(), 0.05, 0.9, 0)
n = 1000
batch_size = 50
base_lr = 0.05
# we need to stop grad of global value to prevent memory leak
lr = f32(base_lr).name("lr").stop_grad()
def get_data(n):
for i in range(n):
x = np.random.rand(batch_size, 1)
y = x*x
yield jt.float32(x), jt.float32(y)
for i,(x,y) in enumerate(get_data(n)):
pred_y = model(x).name("pred_y")
loss = ((pred_y - y)**f32(2)).name("loss")
loss_mean = loss.mean()
SGD.step(loss_mean)
if i>2:
assert prev == jt.liveness_info(), f"memory leak {prev} {jt.liveness_info()}"
prev = jt.liveness_info()
if (i % 10 == 9):
print(f"step {i}, loss = {loss_mean.data.sum()} {jt.liveness_info()}")
else:
loss_mean.data.sum()
jt.liveness_info()
# result is 0.00038617782411165535
result = 0.00038617782411165535
assert abs(loss_mean.data - result) < 1e-6, [loss_mean.data, result]
jt.clean()
def test_backward(self):
np.random.seed(0)
jt.set_seed(3)
model = Model()
SGD = jt.nn.SGD(model.parameters(), 0.05, 0.9, 0)
n = 1000
batch_size = 50
base_lr = 0.05
# we need to stop grad of global value to prevent memory leak
lr = f32(base_lr).name("lr").stop_grad()
def get_data(n):
for i in range(n):
x = np.random.rand(batch_size, 1)
y = x*x
yield jt.float32(x), jt.float32(y)
for i,(x,y) in enumerate(get_data(n)):
pred_y = model(x).name("pred_y")
loss = ((pred_y - y)**f32(2)).name("loss")
loss_mean = loss.mean()
SGD.step(loss_mean)
if i>2:
assert prev == jt.liveness_info(), f"memory leak {prev} {jt.liveness_info()}"
prev = jt.liveness_info()
if (i % 10 == 9):
print(f"step {i}, loss = {loss_mean.data.sum()} {jt.liveness_info()}")
else:
loss_mean.data.sum()
jt.liveness_info()
possible_results = [0.00022486248053610325, 0.00020916158973705024]
loss_mean = loss_mean.data
assert any(abs(loss_mean - r) < 1e-6 for r in possible_results)
jt.clean()
def test_data(self):
a = jt.array([1, 2, 3])
assert (a.data == [1, 2, 3]).all()
d = a.data
a.data[1] = -2
assert (a.data == [1, -2, 3]).all()
assert (a.fetch_sync() == [1, -2, 3]).all()
li = jt.liveness_info()
del a
assert li == jt.liveness_info()
del d
assert li != jt.liveness_info()
def test_zmem_leak3(self):
def test():
class MyFunc(Function):
def execute(self, x, z, y):
self.x = x
self.y = y
return x * y, "test", x / y
def grad(self, grad0, _, grad1):
assert _ is None
res = (grad0 * self.y, None, grad1 * self.x)
return res
a = jt.array(3.0)
b = jt.array(4.0)
c, _, d = MyFunc()(a, "a", b)
g = jt.grad(c + d * 3, [a, b])
jt.sync(g)
import resource
t1 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
for i in range(100000):
test()
if i % 10000 == 0:
jt.clean()
t2 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
for i in range(1000000):
test()
if i % 10000 == 0:
jt.clean()
t3 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
print(t1, t2, t3)
assert t3 < t2 + 10, (t1, t2, t3)
self.assertEqual(jt.liveness_info()["lived_vars"], 0)
def test_zmem_leak(self):
def test():
self.test_multi_grads_multi_out5()
test()
jt.clean()
self.assertEqual(jt.liveness_info()["lived_vars"], 0)
def test_no_grad(self):
a = jt.array(1.0)
with jt.no_grad():
b = a
for i in range(10):
b = b.clone() + 1
assert b.data == 11
jt.clean()
assert jt.liveness_info()["lived_vars"] == 2
def test1(self):
np.random.seed(0)
jt.set_seed(3)
n = 1000
batch_size = 50
base_lr = 0.05
# we need to stop grad of global value to prevent memory leak
lr = f32(base_lr).name("lr").stop_grad()
def get_data(n):
for i in range(n):
x = np.random.rand(batch_size, 1)
y = x * x
yield jt.float32(x), jt.float32(y)
model = Model(input_size=1)
ps = model.parameters()
for i, (x, y) in enumerate(get_data(n)):
pred_y = model(x).name("pred_y")
loss = ((pred_y - y)**f32(2)).name("loss")
loss_mean = loss.mean()
gs = jt.grad(loss_mean, ps)
for p, g in zip(ps, gs):
p -= g * lr
if i > 2:
assert prev == jt.liveness_info(
), f"memory leak {prev} {jt.liveness_info()}"
prev = jt.liveness_info()
print(
f"step {i}, loss = {loss_mean.data.sum()} {jt.liveness_info()}"
)
possible_results = [
0.0009948202641680837,
0.001381353591568768,
0.00110957445576787,
]
loss_mean = loss_mean.data
assert any(abs(loss_mean - r) < 1e-6 for r in possible_results)
jt.clean()
def test1(self):
np.random.seed(0)
jt.set_seed(3)
n = 1000
batch_size = 50
base_lr = 0.05
# tune accumulation_steps for step and batch_size
accumulation_steps = 10
n *= accumulation_steps
batch_size //= accumulation_steps
# we need to stop grad of global value to prevent memory leak
lr = f32(base_lr).name("lr").stop_grad()
def get_data(n):
for i in range(n):
x = np.random.rand(batch_size, 1)
y = x*x
yield jt.float32(x), jt.float32(y)
model = Model(input_size=1)
ps = model.parameters()
opt = Optimizer(ps, lr)
all_loss = 0
for i,(x,y) in enumerate(get_data(n)):
pred_y = model(x).name("pred_y")
loss = ((pred_y - y)**f32(2)).name("loss")
loss_mean = loss.mean() / accumulation_steps
all_loss += loss_mean.item()
opt.backward(loss_mean)
if (i+1) % accumulation_steps == 0:
opt.step()
if i>50:
assert prev == jt.liveness_info(), f"memory leak {prev} {jt.liveness_info()}"
prev = jt.liveness_info()
print(f"step {i}, loss = {loss_mean.data.sum()} {jt.liveness_info()}")
print(all_loss)
possible_results = [19.8639366890402, 8.207454475712439]
assert any(abs(all_loss - r) < 1e-3 for r in possible_results)
jt.clean()
def test_backward_cuda(self):
with jt.flag_scope(use_cuda=1):
np.random.seed(0)
jt.set_seed(3)
model = Model()
SGD = jt.nn.SGD(model.parameters(), 0.05, 0.9, 0)
n = 1000
batch_size = 50
base_lr = 0.05
# we need to stop grad of global value to prevent memory leak
lr = f32(base_lr).name("lr").stop_grad()
def get_data(n):
for i in range(n):
x = np.random.rand(batch_size, 1)
y = x * x
yield jt.float32(x), jt.float32(y)
for i, (x, y) in enumerate(get_data(n)):
pred_y = model(x).name("pred_y")
# cuda x**2.0 will return nan
loss = ((pred_y - y).sqr()).name("loss")
loss_mean = loss.mean()
SGD.step(loss_mean)
if i > 2:
assert prev == jt.liveness_info(
), f"memory leak {prev} {jt.liveness_info()}"
prev = jt.liveness_info()
if (i % 10 == 9):
print(
f"step {i}, loss = {loss_mean.data.sum()} {jt.liveness_info()}"
)
else:
loss_mean.data.sum()
jt.liveness_info()
# result is 0.00018236637697555125
result = 0.00018236637697555125
assert abs(loss_mean.data - result) < 1e-2
jt.clean()
def test_zmem_leak2(self):
def test():
class MyFunc(Function):
def execute(self, x, z, y):
self.x = x.name("x")
self.y = y.name("y")
return x*y, "test", x/y
def grad(self, grad0, _, grad1):
assert _ is None
res = (grad0 * self.y, None, grad1 * self.x)
return res
a = jt.array(3.0).name('a')
b = jt.array(4.0).name('b')
c,_,d = MyFunc()(a, "a", b)
c.name('c'), d.name('d')
g = jt.grad(c+d*3, [a, b])
test()
jt.clean()
self.assertEqual(jt.liveness_info()["lived_vars"], 0)
def run():
start_time = time.time()
fop_num = 10000
fop_input_num = (2, 3) # (i,j) -> [i,i+j] -> [2, 5]
# fop_output_num = (1, 0) # [1,1]
inner_op_num = (0, 3)
fop_type_num = 63 # how many different fuse op
input_queue_num = 15
queue = [1.0]*(input_queue_num+1)
x = get_xorshf96()
rand = lambda x, l, r: l+((x())&r)
ops = ["add", "subtract", "multiply", "divide"]
get_op = lambda x: ops[(x())&3]
for i in range(fop_num):
prev = bc(queue[rand(x,0,input_queue_num)])
y = get_xorshf96(x()&fop_type_num)
inum = rand(y, *fop_input_num)
q = [prev]
for i in range(inum-1):
n = bc(queue[rand(x,0,input_queue_num)])
prev = jt.binary(prev, n, get_op(y))
q.append(prev)
innum = rand(y,*inner_op_num)
for _ in range(innum):
j = rand(y,0,len(q)-1)
n = q[j]
prev = jt.binary(prev, n, get_op(y))
q[j] = prev
prev = rd(prev)
queue[rand(x,0,input_queue_num)] = prev
a = jt.array(0.0)
for x in queue:
a += x
LOG.i("build graph", time.time()-start_time, jt.liveness_info().values())
start_time = time.time()
a.sync()
LOG.i("execute", time.time()-start_time)
batch_size = 50
base_lr = 0.05
# we need to stop grad of global value to prevent memory leak
lr = f32(base_lr).name("lr").stop_grad()
def get_data(n):
for i in range(n):
x = np.random.rand(batch_size, 1)
y = x * x
yield np.float32(x), np.float32(y)
for i, (x, y) in enumerate(get_data(n)):
pred_y = model(x).name("pred_y")
loss = ((pred_y - y)**f32(2)).name("loss")
loss_mean = loss.mean()
ps = jt.find_vars('model')
gs = jt.grad(loss_mean, ps)
for p, g in zip(ps, gs):
p -= g * lr
if i > 2:
assert prev == jt.liveness_info(
), f"memory leak {prev} {jt.liveness_info()}"
prev = jt.liveness_info()
print(f"step {i}, loss = {loss_mean().sum()}")
# result is 0.0009948202641680837
result = 0.0009948202641680837
assert abs(loss_mean.data - result) < 1e-6
