def __enter__(self):
if self._block._empty_prefix:
return self
self._old_scope = getattr(_BlockScope._current, "value", None)
_BlockScope._current.value = self
self._name_scope = _name.Prefix(self._block.prefix)
self._name_scope.__enter__()
self._profiler_scope = _profiler.Scope(
self._block._profiler_scope_name)
self._profiler_scope.__enter__()
return self
def test_gpu_memory_profiler_symbolic():
iter_num = 5
enable_profiler('test_profiler.json', False, False)
profiler.set_state('run')
with profiler.Scope("tensordot"):
A = mx.sym.Variable('A')
B = mx.sym.Variable('B')
C = mx.symbol.dot(A, B, name='dot')
executor = C.simple_bind(mx.gpu(), 'write', A=(4096, 4096), B=(4096, 4096))
a = mx.random.uniform(-1.0, 1.0, shape=(4096, 4096))
b = mx.random.uniform(-1.0, 1.0, shape=(4096, 4096))
a.copyto(executor.arg_dict['A'])
b.copyto(executor.arg_dict['B'])
for i in range(iter_num):
executor.forward()
c = executor.outputs[0]
mx.nd.waitall()
profiler.set_state('stop')
profiler.dump(True)
expected_alloc_entries = [{
'Attribute Name': 'tensordot:in_arg:A',
'Requested Size': str(4 * a.size)
}, {
'Attribute Name': 'tensordot:in_arg:B',
'Requested Size': str(4 * b.size)
}, {
'Attribute Name': 'tensordot:arg_grad:A',
'Requested Size': str(4 * a.size)
}, {
'Attribute Name': 'tensordot:arg_grad:B',
'Requested Size': str(4 * b.size)
}, {
'Attribute Name': 'tensordot:dot',
'Requested Size': str(4 * c.size)
}, {
'Attribute Name': 'tensordot:dot_head_grad',
'Requested Size': str(4 * c.size)
}]
# Sample gpu_memory_profile.csv:
# "Attribute Name","Requested Size","Device","Actual Size","Reuse?"
# "tensordot:arg_grad:A","67108864","0","67108864","0"
# "tensordot:arg_grad:B","67108864","0","67108864","0"
# "tensordot:dot","67108864","0","67108864","0"
# "tensordot:dot_head_grad","67108864","0","67108864","0"
# "tensordot:in_arg:A","67108864","0","67108864","0"
# "tensordot:in_arg:B","67108864","0","67108864","0"
with open('gpu_memory_profile-pid_%d.csv' % (os.getpid()),
mode='r') as csv_file:
csv_reader = csv.DictReader(csv_file)
for expected_alloc_entry in expected_alloc_entries:
csv_file.seek(0)
entry_found = False
for row in csv_reader:
if row['Attribute Name'] == expected_alloc_entry[
'Attribute Name']:
assert row['Requested Size'] == expected_alloc_entry['Requested Size'], \
"requested size={} is not equal to the expected size={}" \
.format(row['Requested Size'],
expected_alloc_entry['Requested Size'])
entry_found = True
break
assert entry_found, \
"Entry for attr_name={} has not been found" \
.format(expected_alloc_entry['Attribute Name'])