Everybody loves %timeit, there’s no doubt about it. So why not have something like that, but for measuring how much memory your line takes? Well, now you can; grab a hold of the script in the following gist and run it like in the example.
# Author: Vlad Niculae <vlad@vene.ro>
# Makes use of memory_profiler from Fabian Pedregosa
# available at https://github.com/fabianp/memory_profiler
from IPython.core.magic import Magics, line_magic, magics_class
class MemMagics(Magics):
@line_magic
def memit(self, line='', setup='pass'):
"""Measure memory usage of a Python statement
Usage, in line mode:
%memit [-ir<R>t<T>] statement
Options:
-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 `np.dot`.
-t<T>: timeout after <T> seconds. Unused if `-i` is active.
Default: None
Examples
--------
::
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,
strict=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:
try:
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):
self.append(x)
q = ListWithPut()
print ('WARNING: cannot import module `multiprocessing`. Forcing '
'the `-i` option.')
run_in_place = True
ns = self.shell.user_ns
def _get_usage(q, stmt, setup='pass', ns={}):
from memory_profiler import memory_usage as _mu
try:
exec setup in ns
_mu0 = _mu()[0]
exec stmt in ns
_mu1 = _mu()[0]
q.put(_mu1 - _mu0)
except Exception as e:
q.put(float('-inf'))
raise e
if run_in_place:
for _ in xrange(repeat):
_get_usage(q, stmt, ns=ns)
else:
# run in consecutive subprocesses
at_least_one_worked = False
for _ in xrange(repeat):
p = pr.Process(target=_get_usage, args=(q, stmt, 'pass', ns))
p.start()
p.join(timeout=timeout)
if p.exitcode == 0:
at_least_one_worked = True
else:
p.terminate()
if p.exitcode == None:
print 'Subprocess timed out.'
else:
print 'Subprocess exited with code %d.' % p.exitcode
q.put(float('-inf'))
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()
ip.register_magics(MemMagics)
Instead of taking care of the dirty process inspection stuff myself, I decided to delegate this to Fabian’s simple but very good memory_profilerGuppy available, but its design seems a bit and overkill for this task.
Please contact me if you find problems with this implementation, this is a preliminary, quick hack-y version.
I know why it’s important to pick the minimum result for timeit, but I’m curious why this is also true for memory.
I didn’t give this too much thought. Maybe it makes more sense to take the maximum. The results are nearly identical over iterations anyway; it should work fine with -r1.