def testPeakMemoryUsage(self, device_type): device = f'{device_type}:0' with ops.device(device): x1 = array_ops.zeros((1000, 1000)) peak1 = config.get_memory_info(device)['peak'] self.assertGreaterEqual(peak1, 4 * 1000 * 1000) with ops.device(device): x2 = array_ops.ones((1000, 1000)) peak2 = config.get_memory_info(device)['peak'] self.assertGreaterEqual(peak2, peak1 + 4 * 1000 * 1000) del x1, x2 # With CPython, causes tensor memory to be immediately freed peak3 = config.get_memory_info(device)['peak'] self.assertGreaterEqual(peak3, peak2) self.assertGreaterEqual(peak3, config.get_memory_info(device)['current'])
def testGetMemoryInfoBasic(self): device = array_ops.zeros([]).backing_device info = config.get_memory_info(device) self.assertGreater(info['current'], 0) self.assertGreater(info['peak'], 0) self.assertEqual(info.keys(), {'current', 'peak'}) self.assertEqual(config.get_memory_usage(device), info['current'])
def testGetMemoryInfoUnknownDevice(self): with self.assertRaisesRegex(ValueError, 'No matching devices found'): config.get_memory_info('unknown_device:0') with self.assertRaisesRegex(ValueError, 'No matching devices found'): config.get_memory_usage('unknown_device:0')
def testGetMemoryUsageSubstring(self, device_type): info = config.get_memory_info(f'{device_type}:0') self.assertGreater(info['current'], 0)
def testGetMemoryInfoCPU(self): with self.assertRaisesRegex(ValueError, 'Allocator stats not available'): config.get_memory_info('CPU:0') with self.assertRaisesRegex(ValueError, 'Allocator stats not available'): config.get_memory_usage('CPU:0')
def testGetMemoryUsageSubstring(self): info = config.get_memory_info('GPU:0') self.assertGreater(info['current'], 0)
def testGetMemoryInfoUnknownDevice(self): with self.assertRaisesRegex(ValueError, 'Failed parsing device name'): config.get_memory_info('unknown_device') with self.assertRaisesRegex(ValueError, 'Failed parsing device name'): config.get_memory_usage('unknown_device')
def testGetMemoryInfoCPU(self): with self.assertRaisesRegex(ValueError, 'CPU does not support'): config.get_memory_info('CPU:0') with self.assertRaisesRegex(ValueError, 'CPU does not support'): config.get_memory_usage('CPU:0')