Decreasing autograd memory usage

[alexbw]

I don't mean "memory leak" in terms of unreachable memory after the Python process quits, I mean memory that is being allocated in the backwards pass, when it should be being freed. I mentioned this problem in #199 , but I think it should be opened as an issue.

For a simple function

import autograd.numpy as np
from autograd import grad

def F(x,z):
    for i in range(100):
        z = np.dot(x,z)
    return np.sum(z)
dF = grad(F)

and a procedure to measure memory usage

from memory_profiler import memory_usage
def make_data():
    np.random.seed(0)
    D = 1000
    x = np.random.randn(D,D)
    x = np.dot(x,x.T)
    z = np.random.randn(D,D)
    return x,z

def m():
    from time import sleep
    x,z = make_data()
    gx = dF(x,z)
    sleep(0.1)
    return gx

mem_usage = np.array(memory_usage(m,interval=0.01))
mem_usage -= mem_usage[0]

and a manual gradient of the same function

def grad_dot_A(g,A,B):
    ga = np.dot(g,B.T)
    ga = np.reshape(ga,np.shape(A))
    return ga

def grad_dot_B(g,A,B):
    gb = np.dot(A.T,g)
    gb = np.reshape(gb, np.shape(B))
    return gb

def dF(x, z):
    z_stack = []
    for i in list(range(100)):
        z_stack.append(z)
        z = np.dot(x, z)
    retval = np.sum(z)

    # Begin Backward Pass
    g_retval = 1
    g_x = 0

    # Reverse of: retval = np.sum(z)
    g_z = repeat_to_match_shape(g_retval, z)
    for i in reversed(list(range(100))):

        # Reverse of: z = np.dot(x, z)
        z = z_stack.pop()
        tmp_g0 = grad_dot_A(g_z, x, z)
        tmp_g1 = grad_dot_B(g_z, x, z)
        g_z = 0
        g_x += tmp_g0
        g_z += tmp_g1
    return g_x

I get the following memory usage profile:

If I replace the dot gradient with the ones used in the manual code, I get the same memory profile, nothing improves.

If I replace the dot product with element-wise multiply, I get a different memory profile, but still not what I would expect:

I would love to help figure this out, but I'm not sure where to start. First thing is of course to document the problem.

[mattjj]

Are you sure there's a leak here? Neither version seems to end at exactly zero net memory usage in the plots (and indeed the manual ones go negative at first), and we would expect autograd to require more memory because it has to trace the forward computation (not just store the values, but box each of them, etc.). This doesn't seem out of line with my expectations, and Python memory deallocation and measurement can be pretty complicated [1] [2].

[mattjj]

Sorry, I don't understand. Doesn't autograd end with less net memory usage than pytorch in that plot? Where is the memory leak shown?

[alexbw]

Updated image of memory usages.

[alexbw]

Here's a gist on a more "real-world" benchmark, taking the gradients of an RNN.

https://gist.github.com/alexbw/7b2a0682f65dd1bcb7120ca2d47a2823

Here's the memory usage:

[alexbw]

For the RNN benchmark, here's separating the forwards and backwards pass of autograd, versus just calling grad(f)(*args).

(apologies, benchmark code is getting a little hairy)
https://gist.github.com/alexbw/5730d3a7f61d36ff482e17b1024ec105

[mattjj]

It's probably from the outgrads being accumulated in a list and summed all at once rather than incrementally summed.

[mattjj]

I just pushed a commit that does in-place gradient accumulation; maybe check how the memory looks with that (though I might not have done it right).

[alexbw]

Much better with inplace-grad-accumulation, now same order of magnitude.

[mattjj]

Btw, thanks for these benchmarks! The latter ones are quite stark. (At first I was just focusing on the 'leak' idea, but not enough on the total memory usage issue.)

[alexbw]

For clarity, just PyTorch v Autograd, so that we don't have to think about anything other than OO AD systems.

PyTorch being a shorter line, btw, means the program is faster (here 2x)

[mattjj]

Is there a potential factor of 2 because of float32 vs float64, or are all these propagating the same dtypes throughout?

I think in-place accumulation makes a lot of sense for computation graphs with high fan-out of some nodes (like an RNN that reuses the same parameters many times; great test!). Fortunately @dougalm's vector space type system made the initial implementation a breeze!

[alexbw]

If I switch everything to float64, I don't see anything appreciably different.

All I did is change all the input and parameter types to float64, that might not be sufficient, though?