Memory issue?

I’ve run into an issue with large matrices and memory. There seem to be two problems:

1. Memory isn’t being released on successive calls of grad. e.g.

a = 10000
b = 10000
A = np.random.randn(a)
B = np.random.randn(b)

def fn(x):
    M = A[:, na] + x[na, :]
    return M[0, 0]

g = grad(fn)

for i in range(100):
    g(B)

is ramping up memory on each iteration.

1. Memory isn’t being released during the backwards pass e.g.

k = 10
def fn(x):
    res = 0
    for i in range(k):
        res = res + np.sum(x)
    return res
g = grad(fn)
b = 200000
g(np.random.randn(b))

This seems to scale in memory (for each call) as O(k), which don’t think is the desired behaviour. For b=150000 this effect does not happen, however.

[mattjj]

We just did some quick experiments and it seems that calling import gc; gc.collect() explicitly does seem to prevent prevent memory from building up.

Also, we think we identified the problem: originally, these lines in backward_pass

while tape:
    node = tape.pop()
    ...

ensured that there were no more circular dependencies when the backward pass finished by popping things off the tape, so that reference counting would be sufficient to clean up garbage. However, to support efficient jacobians I added this line:

tape = copy.copy(tape)  # <-- problem!
while tape:
    node = tape.pop()
    ...

which had the unintended result of preserving circular dependencies and hence preventing reference counting cleanup.

Still, after tape goes out of scope (when grad returns), gc.collect() should be able to garbage collect these things, and so adding explicit calls to gc.collect() to user-level code after calling grad should work. (We could also add it to the definition of grad.)

We’re going to implement a fix (and probably eliminate circular references with tapes in the long term), but for now if you’re having memory problems try calling gc.collect() in your code after computing a gradient.

[jackkamm]

I’ve been using the WeakKeyDict without issue for a few months now, ever since it was first proposed earlier in this thread. Also, I tried py.test in autograd/test with and without WeakKeyDict a few days ago, the output of the unit tests was the same.

But I am happy with your solution – WeakKeyDict is bit too mysterious/magical for my tastes, and explicitly doing the dereferencing seems less likely to cause errors in the future.