Create a gist now

Instantly share code, notes, and snippets.

memit: magic memory usage benching for IPython
# Author: Vlad Niculae <>
# Makes use of memory_profiler from Fabian Pedregosa
# available at
from IPython.core.magic import Magics, line_magic, magics_class
class MemMagics(Magics):
def memit(self, line='', setup='pass'):
"""Measure memory usage of a Python statement
Usage, in line mode:
%memit [-ir<R>t<T>] statement
-r<R>: repeat the loop iteration <R> times and take the best result.
Default: 3
-i: run the code in the current environment, without forking a new
process. This is required on some MacOS versions of Accelerate if your
line contains a call to ``.
-t<T>: timeout after <T> seconds. Unused if `-i` is active.
Default: None
In [1]: import numpy as np
In [2]: %memit np.zeros(1e7)
maximum of 3: 76.402344 MB per loop
In [3]: %memit np.ones(1e6)
maximum of 3: 7.820312 MB per loop
In [4]: %memit -r 10 np.empty(1e8)
maximum of 10: 0.101562 MB per loop
In [5]: memit -t 3 while True: pass;
Subprocess timed out.
Subprocess timed out.
Subprocess timed out.
ERROR: all subprocesses exited unsuccessfully. Try again with the
`-i` option.
maximum of 3: -inf MB per loop
opts, stmt = self.parse_options(line, 'r:t:i', posix=False,
repeat = int(getattr(opts, 'r', 3))
if repeat < 1:
repeat == 1
timeout = int(getattr(opts, 't', 0))
if timeout <= 0:
timeout = None
run_in_place = hasattr(opts, 'i')
# Don't depend on multiprocessing:
import multiprocessing as pr
from multiprocessing.queues import SimpleQueue
q = SimpleQueue()
except ImportError:
class ListWithPut(list):
"Just a list where the `append` method is aliased to `put`."
def put(self, x):
q = ListWithPut()
print ('WARNING: cannot import module `multiprocessing`. Forcing '
'the `-i` option.')
run_in_place = True
ns =
def _get_usage(q, stmt, setup='pass', ns={}):
from memory_profiler import memory_usage as _mu
exec setup in ns
_mu0 = _mu()[0]
exec stmt in ns
_mu1 = _mu()[0]
q.put(_mu1 - _mu0)
except Exception as e:
raise e
if run_in_place:
for _ in xrange(repeat):
_get_usage(q, stmt, ns=ns)
# run in consecutive subprocesses
at_least_one_worked = False
for _ in xrange(repeat):
p = pr.Process(target=_get_usage, args=(q, stmt, 'pass', ns))
if p.exitcode == 0:
at_least_one_worked = True
if p.exitcode == None:
print 'Subprocess timed out.'
print 'Subprocess exited with code %d.' % p.exitcode
if not at_least_one_worked:
print ('ERROR: all subprocesses exited unsuccessfully. Try '
'again with the `-i` option.')
usages = [q.get() for _ in xrange(repeat)]
usage = max(usages)
print u"maximum of %d: %f MB per loop" % (repeat, usage)
if __name__ == '__main__':
ip = get_ipython()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment