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