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