How to make memory efficient operations?

AleksandraDeis: <p>Hi! I’m playing with custom activations and I used the following function instead of ReLU. After doing this I’m running out of GPU memory on Cuda.</p>
<pre><code class="lang-auto">def custom_activation(x):
    return x * tanh(softplus(x)) = x * tanh(ln(1 + exp(x)))
</code></pre>
<p>How can I make this code more memory efficient? May be you can recommend some in-place operations which I can use here?<br>
Or may be there are more efficient ways to implement custom activations? Or may be there are some ways to manage caching on GPU?</p>

ptrblck: <p>Was the code running fine before and raises an out of memory error, if you use these methods?</p>

AleksandraDeis: <p>I am sure that these methods cause the error, because this error is raised when I replace relu activations to this one.<br>
There is even more. If I use this implementation of swish I run out of memory:</p>
<pre><code class="lang-auto">def swish(input):
    return input * torch.sigmoid(input)
</code></pre>
<p>But if I change it as follows everything fits into memory:</p>
<pre><code class="lang-auto">def swish(input, inplace = False):
input = input * torch.sigmoid_(input)
    return input
</code></pre>
<p>So I am looking for the ways to optimize tanh and softmax and make those more memory efficient.</p>

ptrblck: <p>The peak memory usage will be higher in the first method, as intermediate tensors are created.<br>
However, the result will <strong>not</strong> be the same!<br>
Have a look at this code snippet:</p>
<pre><code class="lang-python">def swish(input):
    return input * torch.sigmoid(input)

def swish_inplace(input):
    input = input * torch.sigmoid_(input)
    return input


x = torch.randn(1000, 1000, device='cuda')

torch.cuda.synchronize()
print('Start')
print(torch.cuda.max_memory_allocated() / 1024**2)
print(torch.cuda.memory_allocated() / 1024**2)

output = swish(x)
# output = swish_inplace(x)

torch.cuda.synchronize()
print('After swish')
print(torch.cuda.max_memory_allocated() / 1024**2)
print(torch.cuda.memory_allocated() / 1024**2)
</code></pre>
<p>For the first run we’ll get:</p>
<pre><code class="lang-python">Start
4.64013671875
4.64013671875
After swish
12.27001953125
8.455078125
</code></pre>
<p>Let’s simplify the numbers a bit and claim our tensor uses ~4MB.<br>
As you can see, the memory usage after the function call will be ~8MB, which is fine, since we are still holding to <code>x</code> and created the return tensor <code>output</code>.<br>
However, you can also see that the peak memory usage was actually ~12MB, which would explain the intermediate tensor of <code>torch.sigmoid(input)</code>.</p>
<p>Now let’s swap the vanilla <code>swish</code> for its inplace method and run it again:</p>
<pre><code class="lang-python">Start
4.64013671875
4.64013671875
After swish
8.455078125
8.455078125
</code></pre>
<p>Awesome! It looks like we are indeed saving some peak memory so we should be using the inplace method, shouldn’t we?<br>
Well, let’s have a look at the results first:</p>
<pre><code class="lang-python">x = torch.randn(1000, 1000, device='cuda')
output = swish(x)
# Clone x to keep the original values for comparison reasons
output_ = swish_inplace(x.clone())

# Compare outputs
print((output == (x * torch.sigmoid(x))).all())
&gt; tensor(1, device='cuda:0', dtype=torch.uint8)
print((output_ == (torch.sigmoid(x) * torch.sigmoid(x))).all())
&gt; tensor(1, device='cuda:0', dtype=torch.uint8)
</code></pre>
<p>Ohoh! The inplace method manipulates <code>x</code> inplace, so that  we are in fact calculating <code>torch.sigmoid(x)*torch.sigmoid(x)</code>, which is wrong!</p>
<p>While inplace methods might save some peak memory usage, I’m always careful about using them and always check the result.</p>

AleksandraDeis: <p>Thanks a lot! Inplace methods are really tricky.<br>
Could you suggest some techniques to save some memory for tanh and softmax operations?</p>