Inplace relu doesn't work for custom autograd functions

username: zzd1992
cooked: <p>When we use <code>nn.ReLU(inplace=True)</code>, it should not improve the GPU memory for back-propagation. This is true for convolutional neural networks.</p>
<p>However, the inplace mode doesn’t work for custom functions.<br>
Specifically, I build a custom autograd function and a custom module based on it.<br>
When I combine this module and inplace relu, it seems that <code>nn.ReLU(inplace=True)</code> doesn’t save memory for BP. Here is an example:</p>
<pre><code class="lang-python">net_1 = nn.Sequential()
net_2 = nn.Sequential()
for i in range(50):
    net_1.add_module("conv_{}".format(i), MyConv2d(32, 32, 3))
    net_2.add_module("conv_{}".format(i), MyConv2d(32, 32, 3))
    net_1.add_module("relu_{}".format(i), nn.ReLU(True))
    net_2.add_module("relu_{}".format(i), nn.ReLU(False))
</code></pre>
<p>Then I get the <code>max_memory_allocated</code> for <code>net_1</code> and <code>net_2</code> given the same input.<br>
And I find the used memories are the same.<br>
net_1: 470MB<br>
net_2: 470MB<br>
How can I modify my custom function to make <code>nn.ReLU(inplace=True)</code> works?</p>

username: gphilip
cooked: <p>Did you try the same code with a built-in CNN instead of <code>MyConv2d</code>? Does that report a significant difference in <code>max_memory_allocated</code>?</p>

username: zzd1992
cooked: <p><a class="mention" href="/u/gphilip">@gphilip</a><br>
If replace <code>MyConv2d</code> with the build-in CNN, inplace mode can save half of the memory.</p>
<p>net_1: 240 MB<br>
net_2: 470 MB</p>

username: zzd1992
cooked: <p>Can someone point out the reason for this?</p>

username: ptrblck
cooked: <p>Could you post an executable code snippet to reproduce the issue, please?</p>

username: zzd1992
cooked: <p><a class="mention" href="/u/gphilip">@gphilip</a><br>
<a class="mention" href="/u/ptrblck">@ptrblck</a><br>
This is the executable code:</p>
<pre><code class="lang-python">import torch
from torch import nn

class conv(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        y = 2 * x
        ctx.save_for_backward(y)
        return torch.sin(y)
    
    @staticmethod
    def backward(ctx, grad_out):
        y = ctx.saved_tensors
        with torch.no_grad():
            if ctx.needs_input_grad[0]:
                return grad_out * torch.cos(y)
            else:
                return None
    

class MyConv(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.parameter.Parameter(torch.randn(1, 1, 64, 64))
        self.main = conv.apply

    def forward(self, x):
        return self.main(x) * self.weight


def evaluate_memory(module):
    x = torch.randn(32, 32, 64, 64).cuda()
    torch.cuda.reset_peak_memory_stats()
    y = module(x)
    return torch.cuda.max_memory_allocated() / (1024 ** 2)

net_1 = nn.Sequential()
net_2 = nn.Sequential()

use_default_conv = False

for i in range(50):
    if use_default_conv:
        net_1.add_module("conv_{}".format(i), nn.Conv2d(32, 32, 3, padding=1))
        net_2.add_module("conv_{}".format(i), nn.Conv2d(32, 32, 3, padding=1))
    else:
        net_1.add_module("conv_{}".format(i), MyConv())
        net_2.add_module("conv_{}".format(i), MyConv())
    net_1.add_module("relu_{}".format(i), nn.ReLU(True))
    net_2.add_module("relu_{}".format(i), nn.ReLU(False))

net_1.cuda()
net_2.cuda()

m_1 = evaluate_memory(net_1)
m_2 = evaluate_memory(net_2)
print("memory for in-place mode: {:.1f}".format(m_1))
print("memory for out-place mode: {:.1f}".format(m_2))

</code></pre>
<p>when set <code>use_defualt_conv=True</code>, the output is</p>
<pre><code class="lang-auto">memory for in-place mode: 819.6
memory for out-place mode: 1619.6
</code></pre>
<p>when set <code>use_defualt_conv=False</code>, the output is</p>
<pre><code class="lang-auto">memory for in-place mode: 2401.6
memory for out-place mode: 2417.6
</code></pre>
<p>My point is that when set <code>use_defualt_conv=False</code>, the memory of out-place mode should be half of in-place mode. However, it is not.</p>

username: gphilip
cooked: <p>(I think you meant to alert <a class="mention" href="/u/ptrblck">@ptrblck</a> ; he is the one who asked for the code.)</p>
<p>Thank you for the code. After playing with it for a bit I discovered something interesting: If I set <code>use_default_conv=False</code> and change</p>
<pre><code class="lang-auto">for i in range(50):
</code></pre>
<p>to</p>
<pre><code class="lang-auto">for i in range(50000):
</code></pre>
<p>then the run takes much more time to complete (as expected), but the memory use reported <strong>still remains</strong> at <code>48.0</code>.</p>
<p>This suggests to me that the number of floats in <code>MyConv</code> is too small to cause any real gains when using <code>ReLU</code> in place, even when there are a 50000 copies of <code>MyConv</code> in the net. This may be the explanation.</p>

username: zzd1992
cooked: <p><a class="mention" href="/u/gphilip">@gphilip</a></p>
<p>Thanks for your remind.<br>
To verify your guess, I modify <code>MyConv</code>. (See the previous post.)<br>
After modification, the memory usage is proportional to the number of layers and it is larger than the default Conv.</p>
<p>when set <code>use_defualt_conv=True</code>, the output is</p>
<pre><code class="lang-auto">memory for in-place mode: 819.6
memory for out-place mode: 1619.6
</code></pre>
<p>when set <code>use_defualt_conv=False</code>, the output is</p>
<pre><code class="lang-auto">memory for in-place mode: 2401.6
memory for out-place mode: 2417.6
</code></pre>
<p>Those results suggest that in-place mode of Relu doesn’t work for custom autograd functions.<br>
In fact, my real custom function is more complex than the default Conv. I just provide an easy example to explain the problem.</p>

username: googlebot
cooked: <p>your example should be changed to:</p>
<pre><code class="lang-auto">    def forward(ctx, x):
        y = torch.sin(2*x)
        if ctx.needs_input_grad[0]: ctx.save_for_backward(y)
        return y
</code></pre>
<p>as your version keeps both 2<em>x and sin(2</em>x) tensors. I’m not sure about relu’s role.</p>

username: gphilip
cooked: <p>Thanks for the modified code. I am getting different results than you when I run a slightly modified version of your code. My results suggest that <code>ReLU</code> being in-place <strong>does</strong> make a significant difference, even for custom modules.</p>
<p>Here is the modified code that I ran:</p>
<pre><code class="lang-auto">import torch
from torch import nn

class MyConv(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.parameter.Parameter(torch.randn(1, 1, 64, 64))
        

    def forward(self, x):
        return (x * self.weight)


def evaluate_memory(module):
    x = torch.rand(32, 32, 64, 64).cuda()
    torch.cuda.reset_peak_memory_stats()
    y = module(x)
    return torch.cuda.max_memory_allocated(torch.device("cuda")) / (1024 ** 2)

net_1 = nn.Sequential()
net_2 = nn.Sequential()

use_default_conv = True

for i in range(50):
    if use_default_conv:
        net_1.add_module("conv_{}".format(i), nn.Conv2d(32, 32, 3, padding=1))
        net_2.add_module("conv_{}".format(i), nn.Conv2d(32, 32, 3, padding=1))
    else:
        net_1.add_module("conv_{}".format(i), MyConv())
        net_2.add_module("conv_{}".format(i), MyConv())
    net_1.add_module("relu_{}".format(i), nn.ReLU(True))
    net_2.add_module("relu_{}".format(i), nn.ReLU(False))

net_1.cuda()
net_2.cuda()

m_1 = evaluate_memory(net_1)
m_2 = evaluate_memory(net_2)
print("memory for in-place mode: {:.1f}".format(m_1))
print("memory for out-place mode: {:.1f}".format(m_2))

</code></pre>
<p>When I run this with <code>use_default_conv = True</code> I get:</p>
<p>memory for in-place mode: 819.6<br>
memory for out-place mode: 1619.6</p>
<p>When I run this with <code>use_default_conv = False</code> I get:</p>
<p>memory for in-place mode: 817.6<br>
memory for out-place mode: 1617.6</p>
<p>This tells that <code>ReLU</code> being in-place does make the same degree of difference, in both cases.</p>

username: zzd1992
cooked: <p><a class="mention" href="/u/gphilip">@gphilip</a><br>
In your case, <code>custom function</code> is not used.<br>
Here, <code>custom function</code> means custom autograd function in PyTorch.<br>
That is a sub-calss <code>torch.autograd.Function</code>. See <code>conv</code> in my code.</p>

username: gphilip
cooked: <p>Thanks for pointing that out, I missed it altogether!</p>
<p>Edited to add:</p>
<p>If you change the line</p>
<pre><code class="lang-auto">ctx.save_for_backward(y)
</code></pre>
<p>to</p>
<pre><code class="lang-auto">ctx.save_for_backward(x)
</code></pre>
<p>then there is still a significant difference in the two memory usages reported (with and without in-place, respectively). This suggests that while using the in-place <code>ReLU</code> does make a difference, this difference is swamped by the other memory being used (and not freed) by your <code>forward</code> method.</p>