added GPU memory and layer runtime profiling to the model info method

cvmodelz/models/meta_info.py

@@ -6,6 +6,7 @@ from typing import Tuple
 from typing import Callable
 
 pyaml.add_representer(types.FunctionType, lambda cls, func: cls.represent_data(str(func)))
+pyaml.add_representer(types.MethodType, lambda cls, func: cls.represent_data(str(func)))
 
 @dataclass
 class ModelInfo(object):
@@ -22,6 +23,9 @@ class ModelInfo(object):
 
 	prepare_func:               Callable    = lambda x: x
 
+	def prepare(self, foo):
+		pass
+
 	def __str__(self):
 		obj = dict(ModelInfo=self.__dict__)
 		return pyaml.dump(obj, sort_dicts=False, )
@@ -30,3 +34,6 @@ class ModelInfo(object):
 if __name__ == '__main__':
 	info = ModelInfo()
 	print(info)
+
+	info = ModelInfo(prepare_func=info.prepare)
+	print(info)

cvmodelz/utils/__init__.py

@@ -7,6 +7,10 @@ import logging
 from functools import partial
 from tabulate import tabulate
 
+from chainer.function_hooks import TimerHook
+from chainer.function_hooks import CupyMemoryProfileHook
+
+
 def get_attr_from_path(obj, path, *, sep="/"):
 	def getter(o, attr):
 		return
@@ -66,7 +70,15 @@ def print_model_info(model, file=sys.stdout, input_size=None, input_var=None):
 	_print(f"Printing some information about \"{name}\" model")
 	_print(tabulate(rows, tablefmt="fancy_grid"))
 
-	shapes = _get_activation_shapes(model, input_size or default_size, input_var)
+	timer_hook = TimerHook()
+	memory_hook = CupyMemoryProfileHook()
+
+	with timer_hook, memory_hook:
+		shapes = _get_activation_shapes(model, input_size or default_size, input_var)
+
+	timer_hook.print_report()
+	memory_hook.print_report()
+
 	_print("In/Out activation shapes:")
 	_print(tabulate(shapes,
 		headers=["Link name", "Input", "Output"],