The order of ReLU and BatchNorm in resnet50 during backProp

jih332: <p>Hi, I was playing with resnet50 model pre-defined in <code>torchvision.models</code>, I use following code to register hooks</p>
<pre><code class="lang-auto">import torchvison.models as tmodels

def is_supported_instance(module):
    for typed_modules in get_supported_modules().values():
        if isinstance(module, tuple(typed_modules)):
            return True
    return False

def fw_hook(module, input, output):
    print(f'forward hook in module {module.__full_name__}')
    
    def bw_hook(grad):
        print(f'backward hook in module {module.__full_name__}')
        
    handle = output.register_hook(bw_hook)
    module.__bw_hook_handle__ = handle

def register_hook(module, hook_fn):
    if is_supported_instance(module):
        handle = module.register_forward_hook(hook_fn)
        module.__fw_hook_handle__ = handle
    
# tm = test_model()
tm = tmodels.resnet50()

for n, m in tm.named_modules():
    m.__full_name__ = n.replace('.', '/')
    
tm.apply(lambda m: register_hook(m, fw_hook))

input = torch.ones([2, 3, 224, 224], dtype = torch.float, requires_grad = True)
out = tm(input)
out = out.mean()
out.backward()
</code></pre>
<p>The result seems good in forward propagation part, however, in backward prop, I noticed that the order of <code>ReLU</code> and <code>BatchNorm2d</code> was reversed, as shown below</p>
<blockquote>
<p>…<br>
forward hook in module layer4/2/conv2<br>
forward hook in module layer4/2/bn2<br>
forward hook in module layer4/2/relu<br>
forward hook in module layer4/2/conv3<br>
forward hook in module layer4/2/bn3<br>
forward hook in module layer4/2/relu<br>
forward hook in module avgpool<br>
forward hook in module fc<br>
backward hook in module fc<br>
backward hook in module avgpool<br>
backward hook in module layer4/2/bn3<br>
backward hook in module layer4/2/relu<br>
backward hook in module layer4/2/conv3<br>
backward hook in module layer4/2/bn2<br>
backward hook in module layer4/2/relu<br>
backward hook in module layer4/2/conv2<br>
…</p>
</blockquote>
<p>besides, I’ve checked the code for resnet50 in <code>torchvision.models</code>, it seems that <code>ReLU</code> always follows a <code>BatchNorm2d</code> layer. So why did this happen?<br>
Thanks in advance.</p>

Diego: <p>Hi, so I am no expert but I assume that this is an optimization decision. If you think about what the RELU operation (e.g <code>max(0, x)</code>) does to the gradients you will notice that for negative inputs, it kills the gradient (turns it to 0) and for positive inputs, it just lets it flow since dx/dx=1. So if you put the RELU first and there are a lot of negative gradients flowing you will just kill them all and save computation in the backward pass.<br>
Either that or there is a bug in your code and Im just making this up <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=9" title=":sweat_smile:" class="emoji" alt=":sweat_smile:">. Anyways I am not sure about this so maybe someone else can clarify further.</p>

jih332: <p>Hey thanks for the reply. I’ve thought about that too but if it was done for the purpose of reducing computation, shouldn’t it be</p>
<blockquote>
<p>backward hook in module layer4/2/relu<br>
backward hook in module layer4/2/bn3<br>
backward hook in module layer4/2/conv3</p>
</blockquote>
<p>which is the expected order as gradient props backward.<br>
I’ve tried the register hook code on another toy model,</p>
<pre><code class="lang-auto">class test_model(nn.Module):
    def __init__(self):
        super(test_model, self).__init__()
        self.conv1 = nn.Conv2d(3, 8, [3, 3])
        self.bn1 = nn.BatchNorm2d(8)
        self.conv2 = nn.Conv2d(8, 3, [3, 3]) 
        self.bn2 = nn.BatchNorm2d(3) 
        self.relu = nn.ReLU()
        self.avgpool = nn.AdaptiveAvgPool2d(3)
        self.linear = nn.Linear(3, 1)
    
    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.avgpool(out)
        out = self.linear(out)
        return out
</code></pre>
<p>for this model the result is as expected</p>
<blockquote>
<p>forward hook in module conv1<br>
forward hook in module bn1<br>
forward hook in module relu<br>
forward hook in module conv2<br>
forward hook in module bn2<br>
forward hook in module relu<br>
forward hook in module avgpool<br>
forward hook in module linear<br>
backward hook in module linear<br>
backward hook in module avgpool<br>
backward hook in module relu<br>
backward hook in module bn2<br>
backward hook in module conv2<br>
backward hook in module relu<br>
backward hook in module bn1<br>
backward hook in module conv1</p>
</blockquote>
<p>so I really have no idea what’s going on there…</p>

Diego: <aside class="quote no-group quote-modified" data-username="jih332" data-post="3" data-topic="50475">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jih332/48/12441_2.png" class="avatar"> jih332:</div>
<blockquote>
<p>Hey thanks for the reply. I’ve thought about that too but if it was done for the purpose of reducing computation, shouldn’t it be</p>
<blockquote>
<p>backward hook in module layer4/2/relu<br>
backward hook in module layer4/2/bn3<br>
backward hook in module layer4/2/conv3</p>
</blockquote>
</blockquote>
</aside>
<p>You are right. Since it was backward I read the prints from the bottom to the top and got confused <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"></p>

jih332: <p>turns out “it’s because <code>ReLU</code> used in torchvision’s <code>resnet50</code> is an inplace operation, given that the operation += just before the ReLU (exists in both BatchNorm and Linear) is also inplace, autograd engine does some extra checks and operation reordering during backward prop.”<br>
I still don’t have a clue how the reordering is done, but I tried to change the <code>ReLU</code> in my toy model above to inplace version, and found that the order of <code>ReLU</code> and <code>BatchNorm</code> was indeed reversed during backward prop. So I guess it confirms the explanation.</p>