def testMemoryLeakAnonymousVariable(self): if not memory_test_util.memory_profiler_is_available(): self.skipTest("memory_profiler required to run this test") def f(): inputs = Variable(array_ops.zeros([32, 100], dtypes.float32)) del inputs memory_test_util.assert_no_leak(f, num_iters=10000)
def testMemoryLeakInSimpleModelForwardOnly(self): if not memory_test_util.memory_profiler_is_available(): self.skipTest("memory_profiler required to run this test") inputs = array_ops.zeros([32, 100], dtypes.float32) net = SingleLayerNet() def f(): with backprop.GradientTape(): net(inputs) memory_test_util.assert_no_leak(f)
def testMemoryLeakInFunction(self): if not memory_test_util.memory_profiler_is_available(): self.skipTest("memory_profiler required to run this test") def f(): @def_function.function def graph(x): return x * x + x graph(constant_op.constant(42)) memory_test_util.assert_no_leak(f, num_iters=1000, increase_threshold_absolute_mb=30)
def testMemoryLeakInSimpleModelForwardAndBackward(self): if not memory_test_util.memory_profiler_is_available(): self.skipTest("memory_profiler required to run this test") inputs = array_ops.zeros([32, 100], dtypes.float32) net = SingleLayerNet() def f(): with backprop.GradientTape() as tape: result = net(inputs) tape.gradient(result, net.variables) del tape memory_test_util.assert_no_leak(f)
def testMemoryLeakInLocalCopy(self): if not memory_test_util.memory_profiler_is_available(): self.skipTest("memory_profiler required to run this test") remote.connect_to_remote_host(self._cached_server_target) # Run a function locally with the input on a remote worker and ensure we # do not leak a reference to the remote tensor. @def_function.function def local_func(i): return i def func(): with ops.device("job:worker/replica:0/task:0/device:CPU:0"): x = array_ops.zeros([1000, 1000], dtypes.int32) local_func(x) memory_test_util.assert_no_leak( func, num_iters=100, increase_threshold_absolute_mb=50)
def testMemoryLeakInGlobalGradientRegistry(self): # Past leak: b/139819011 if not memory_test_util.memory_profiler_is_available(): self.skipTest("memory_profiler required to run this test") def f(): @def_function.function(autograph=False) def graph(x): @def_function.function(autograph=False) def cubed(a): return a * a * a y = cubed(x) # To ensure deleting the function does not affect the gradient # computation. del cubed return gradient_ops.gradients(gradient_ops.gradients(y, x), x) return graph(constant_op.constant(1.5))[0].numpy() memory_test_util.assert_no_leak(f, num_iters=300, increase_threshold_absolute_mb=50)