Skip to content

Commit

Permalink
[NFC] polish colossalai/utils/tensor_detector/tensor_detector.py code…
Browse files Browse the repository at this point in the history 
… style (#1566)
  • Loading branch information
@Gy-Lu
Gy-Lu committed Sep 8, 2022
1 parent ab546e5 commit b387738
Showing 1 changed file with 32 additions and 36 deletions.
68 changes: 32 additions & 36 deletions colossalai/utils/tensor_detector/tensor_detector.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,18 +5,17 @@
from typing import Optional
from collections import defaultdict


LINE_WIDTH = 108
LINE = '-' * LINE_WIDTH + '\n'


class TensorDetector():

def __init__(self,
show_info: bool = True,
log: str = None,
include_cpu: bool = False,
module: Optional[nn.Module] = None
):
module: Optional[nn.Module] = None):
"""This class is a detector to detect tensor on different devices.
Args:
Expand All @@ -28,7 +27,7 @@ def __init__(self,
"""
self.show_info = show_info
self.log = log
self.include_cpu = include_cpu
self.include_cpu = include_cpu
self.tensor_info = defaultdict(list)
self.saved_tensor_info = defaultdict(list)
self.order = []
Expand Down Expand Up @@ -57,13 +56,13 @@ def get_tensor_mem(self, tensor):

def mem_format(self, real_memory_size):
# format the tensor memory into a reasonal magnitude
if real_memory_size >= 2 ** 30:
return str(real_memory_size / (2 ** 30)) + ' GB'
if real_memory_size >= 2 ** 20:
return str(real_memory_size / (2 ** 20)) + ' MB'
if real_memory_size >= 2 ** 10:
return str(real_memory_size / (2 ** 10)) + ' KB'
return str(real_memory_size) + ' B'
if real_memory_size >= 2**30:
return str(real_memory_size / (2**30)) + ' GB'
if real_memory_size >= 2**20:
return str(real_memory_size / (2**20)) + ' MB'
if real_memory_size >= 2**10:
return str(real_memory_size / (2**10)) + ' KB'
return str(real_memory_size) + ' B'

def collect_tensors_state(self):
for obj in gc.get_objects():
Expand All @@ -74,11 +73,11 @@ def collect_tensors_state(self):
self.detected.append(id(obj))
# skip paramters we had added in __init__ when module is an instance of nn.Module for the first epoch
if id(obj) not in self.tensor_info:

name = type(obj).__name__
# after backward, we want to update the records, to show you the change
if isinstance(self.module, nn.Module) and name == 'Parameter':
if obj.grad is not None:
if obj.grad is not None:
# with grad attached
for par_name, param in self.module.named_parameters():
if param.requires_grad and param.grad.equal(obj.grad):
Expand All @@ -88,7 +87,7 @@ def collect_tensors_state(self):
# there will be no new paramters created during running
# so it must be in saved_tensor_info
continue
# we can also marked common tensors as tensor(with grad)
# we can also marked common tensors as tensor(with grad)
if name == 'Tensor' and (obj.is_leaf or obj.retains_grad):
if obj.grad is not None:
name = name + ' (with grad)'
Expand All @@ -104,7 +103,7 @@ def collect_tensors_state(self):
self.tensor_info[id(obj)].append(obj.dtype)
self.tensor_info[id(obj)].append(self.get_tensor_mem(obj))
# recorded the order we got the tensor
# by this we can guess the tensor easily
# by this we can guess the tensor easily
# it will record every tensor updated this turn
self.order.append(id(obj))
# recorded all different devices
Expand All @@ -114,44 +113,41 @@ def collect_tensors_state(self):
def print_tensors_state(self):
template_format = '{:3s}{:<30s}{:>10s}{:>20s}{:>10s}{:>20s}{:>15s}'
self.info += LINE
self.info += template_format.format(' ', 'Tensor', 'device', 'shape', 'grad', 'dtype', 'Mem')
self.info += template_format.format(' ', 'Tensor', 'device', 'shape', 'grad', 'dtype', 'Mem')
self.info += '\n'
self.info += LINE

# if a tensor updates this turn, and was recorded before
# it should be updated in the saved_tensor_info as well
outdated = [x for x in self.saved_tensor_info.keys() if x in self.order]
minus = [x for x in self.saved_tensor_info.keys() if x not in self.detected]
minus = outdated + minus
minus = outdated + minus
if len(self.order) > 0:
for tensor_id in self.order:
self.info += template_format.format('+',
str(self.tensor_info[tensor_id][0]),
str(self.tensor_info[tensor_id][1]),
str(tuple(self.tensor_info[tensor_id][2])),
str(self.tensor_info[tensor_id][3]),
str(self.tensor_info[tensor_id][4]),
str(self.tensor_info[tensor_id][5]))
self.info += template_format.format('+', str(self.tensor_info[tensor_id][0]),
str(self.tensor_info[tensor_id][1]),
str(tuple(self.tensor_info[tensor_id][2])),
str(self.tensor_info[tensor_id][3]),
str(self.tensor_info[tensor_id][4]),
str(self.tensor_info[tensor_id][5]))
self.info += '\n'
if len(self.order) > 0 and len(minus) > 0:
self.info += '\n'
if len(minus) > 0:
for tensor_id in minus:
self.info += template_format.format('-',
str(self.saved_tensor_info[tensor_id][0]),
str(self.saved_tensor_info[tensor_id][1]),
str(tuple(self.saved_tensor_info[tensor_id][2])),
str(self.saved_tensor_info[tensor_id][3]),
str(self.saved_tensor_info[tensor_id][4]),
str(self.saved_tensor_info[tensor_id][5]))
self.info += template_format.format('-', str(self.saved_tensor_info[tensor_id][0]),
str(self.saved_tensor_info[tensor_id][1]),
str(tuple(self.saved_tensor_info[tensor_id][2])),
str(self.saved_tensor_info[tensor_id][3]),
str(self.saved_tensor_info[tensor_id][4]),
str(self.saved_tensor_info[tensor_id][5]))
self.info += '\n'
# deleted the updated tensor
self.saved_tensor_info.pop(tensor_id)


# trace where is the detect()
locate_info = inspect.stack()[2]
locate_msg = '"' + locate_info.filename + '" line ' + str(locate_info.lineno)
locate_msg = '"' + locate_info.filename + '" line ' + str(locate_info.lineno)

self.info += LINE
self.info += f"Detect Location: {locate_msg}\n"
Expand All @@ -167,8 +163,8 @@ def print_tensors_state(self):
if self.log is not None:
with open(self.log + '.log', 'a') as f:
f.write(self.info)
def detect(self, include_cpu = False):

def detect(self, include_cpu=False):
self.include_cpu = include_cpu
self.collect_tensors_state()
self.print_tensors_state()
Expand All @@ -180,4 +176,4 @@ def detect(self, include_cpu = False):

def close(self):
self.saved_tensor_info.clear()
self.module = None
self.module = None

0 comments on commit b387738

Please sign in to comment.