Source code for golem.utilities.profiler.memory_profiler

import inspect
import os
import time

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

from golem.core.dag.graph import Graph
from golem.core.optimisers.optimizer import GraphOptimizer
from golem.utilities.requirements_notificator import warn_requirement

try:
    import objgraph

    from memory_profiler import memory_usage
except ImportError:
    warn_requirement('objgraph', 'golem[profilers]', should_raise=True)


[docs]class MemoryProfiler: """Visual interpretation of memory usage. Create two ``png`` files. Args: function: function to profile. path: path to save profiling result. list args: arguments for function in array format. dict kwargs: arguments for function in dictionary format. list roots: array with classes used as ROOT nodes in the call-graph. int max_depth: maximum depth of graph. """ def __init__(self, function, path: str, args=None, kwargs=None, roots=None, max_depth: int = 7, visualization=False): if args is None: args = [] if kwargs is None: kwargs = {} self.folder = os.path.abspath(path) if not os.path.exists(self.folder): os.mkdir(self.folder) if roots is None: roots = [Graph, GraphOptimizer] if visualization: # # Create the plot of the memory time dependence. # self._create_memory_plot(function, args, kwargs) # Create call graph. self._create_memory_graph(roots, max_depth) def _create_memory_plot(self, function, args, kwargs, interval: float = 0.1): start_time = time.time() mem_res = memory_usage((function, args, kwargs), interval=interval) total_time = time.time() - start_time length = len(mem_res) division = total_time / length time_split = np.linspace(division, total_time, length) max_index = np.argmax(mem_res) plt.figure(figsize=(12, 8)) sns.set_style("whitegrid") ax = sns.lineplot((time_split, mem_res), linewidth=3, marker="X", markersize=8) ax.axhline(mem_res[max_index], ls='--', color='red') ax.axvline(time_split[max_index], ls='--', color='red') plt.legend(labels=[f'interval {interval}', 'max memory usage']) ax.set_title('Memory time dependence', size=25) ax.set_xlabel("time [s]", fontsize=16) ax.set_ylabel("memory used [KB]", fontsize=16) filename = os.path.join(self.folder, 'memory_plot.png') ax.figure.savefig(filename) def _create_memory_graph(self, roots, max_depth): filename = os.path.join(self.folder, 'memory_graph.png') objgraph.show_refs( roots, max_depth=max_depth, filename=filename, filter=lambda x: not type(x).__name__ in ['module', 'str', 'code', 'IPythonKernel', '_Helper', 'builtin_function_or_method', 'type', '_Printer', 'Printer', 'float64', 'float', 'int'], highlight=lambda x: inspect.isclass(x) or inspect.isfunction(x), refcounts=True )