When debugging, I often want to measure a piece of code's performance/memory characteristics. I do this with the stdlib time class and psutil module.
These are some snippets I always have in my utils before starting a project.
Note: I use perf_counter (i.e performance counter clock) to measure time. For more details, see this.
Measuring time and memory
I use python's contextmanager to manage state. Realpython has a nice post on this.
Using the context manager class
import time
import psutil
class TimeMem():
def __init__(self, step_desc):
self.step_desc = step_desc
@staticmethod
def mem_info():
"""
Returns process memory-usage in bytes
"""
return psutil.Process(os.getpid()).memory_info().rss
def delta(self):
"""
Time taken in seconds
"""
return time.perf_counter() - self.tic
def mem_consumed(self, div=1):
"""
Returns memory consumed when the context was active.
Defaults to bytes.
"""
return (self.mem_info() - self.mem_start) / div
def __enter__(self):
self.tic = time.perf_counter()
self.mem_start = self.mem_info()
logging.info(f"Start: {self.step_desc}")
def __exit__(self, exc_type, exc_value, exc_tb):
logging.info(f"Done {self.step_desc} "
f"Time {self.delta():.2f} seconds "
f"Mem {self.mem_consumed(1024**2):.2f} MB")Using the decorator syntax
The context manager API can be verbose. Using the contextmanager decorator makes it more clear.
from contextlib import contextmanager
class Timer():
def __init__(self):
self.tic = time.perf_counter()
def delta(self):
"""
Time taken in seconds
"""
return time.perf_counter() - self.tic
class Memory():
def __init__(self):
self.initial = Memory.mem_info()
@staticmethod
def mem_info():
"""
Returns process memory-usage in bytes
"""
return psutil.Process(os.getpid()).memory_info().rss
def mem_consumed(self, div=1):
"""
Returns memory consumed when the context was active.
Defaults to bytes.
"""
return (self.mem_info() - self.initial) / div
@contextmanager
def measure(step_name):
t = Timer()
m = Memory()
try:
logging.info(f"Start {step_name}")
yield
finally:
logging.info(f"Done {step_name} "
f"Time {t.delta():.2f} seconds "
f"Mem {m.mem_consumed(1024**2):.2f} MB")import numpy
with(TimeMem("TimMem: Init numpy array")):
result = numpy.arange(10**9, dtype=numpy.int64)
with(measure("decorator: Init numpy array")):
result = numpy.arange(10**9, dtype=numpy.int64)
###
2023-03-11 13:43:44,567 root INFO Start: TimMem: Init numpy array
2023-03-11 13:43:45,541 root INFO Done TimMem: Init numpy array Time 0.97 seconds Mem 7629.34 MB
2023-03-11 13:43:45,542 root INFO Start decorator: Init numpy array
2023-03-11 13:43:46,523 root INFO Done decorator: Init numpy array Time 0.98 seconds Mem 7629.34 MB