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')
