Difference between extracting from forward hook and regular forward pass

Nikolas_McNeal: <p>I am passing an image through CLIP ResNet50 and extracting an intermediate layer. I notice that I get different activation values if I use a forward hook or if I manually create the forward pass (passing through the elements in named_modules()). However, I get the same outputs for conv layers. For batchnorm and ReLU, I get different values. I am not sure why this is happening.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import clip
from PIL import Image
import numpy as np


device = "cuda" if torch.cuda.is_available() else "cpu"
clip_model, preprocess = clip.load("RN50", device=device)

img = np.random.random((224,224,3)).astype(np.float16)
img = Image.fromarray(img, 'RGB')
image_input = preprocess(img).unsqueeze(0).to(device).half()


class ForwardPassVersion(nn.Module):
    def __init__(self, encoder):
        super(ForwardPassVersion, self).__init__()
        self.model = encoder
    
    def forward(self, x):
        x = self.model.visual.conv1(x)
        x = self.model.visual.bn1(x)
        x = self.model.visual.relu1(x)
        x = self.model.visual.conv2(x)
        #x = self.model.visual.bn2(x)
        
        x = torch.flatten(x).to(torch.float32)

        return x

class HookVersion(nn.Module):
    def __init__(self, encoder):
        super(HookVersion, self).__init__()
        self.model = encoder
        self.layer_name = 'conv2'
        self.activation = {}
        
        layer = dict([*self.model.visual.named_modules()])[self.layer_name]
        layer.register_forward_hook(self.get_activation(self.layer_name, self.activation))

    def get_activation(self, name, activations):
        def hook(model, input, output):
            activations[name] = output
        return hook

    def forward(self, x):
        _ = self.model.encode_image(x) 
        x = self.activation[self.layer_name]      
        x = torch.flatten(x).to(torch.float32)
        return x
    

fp_model = ForwardPassVersion(clip_model).to(device)
hook_model = HookVersion(clip_model).to(device)

fp_output = fp_model(image_input)
hook_output = hook_model(image_input)

print(torch.allclose(fp_output, hook_output, atol=1e-4))
</code></pre>
<p>The above code returns “True”. However, if I uncomment the line in ForwardPassVersion and change self.layer_name to “bn2” in HookVersion, it returns “False”.</p>

ptrblck: <p>Check if the output activation is used in any inplace operation afterwards.<br>
If not, could you please post a minimal and executable code snippet reproducing the issue?</p>

Nikolas_McNeal: <p>Yes, I have updated the post with a minimal and executable code snippet reproducing the issue. I don’t believe the output activation is used in any inplace operation after.</p>

ptrblck: <p>Thanks for updating the code!<br>
I see <code>True</code> as the result but assume I should see a failure in <code>torch.allclose</code>?</p>

Nikolas_McNeal: <p>Yes, I expected True in all cases. However, if I change the code to (changes commented):</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import clip
from PIL import Image
import numpy as np


device = "cuda" if torch.cuda.is_available() else "cpu"
clip_model, preprocess = clip.load("RN50", device=device)

img = np.random.random((224,224,3)).astype(np.float16)
img = Image.fromarray(img, 'RGB')
image_input = preprocess(img).unsqueeze(0).to(device).half()


class ForwardPassVersion(nn.Module):
    def __init__(self, encoder):
        super(ForwardPassVersion, self).__init__()
        self.model = encoder
    
    def forward(self, x):
        x = self.model.visual.conv1(x)
        x = self.model.visual.bn1(x)
        x = self.model.visual.relu1(x)
        x = self.model.visual.conv2(x)
        x = self.model.visual.bn2(x) # added this line
        
        x = torch.flatten(x).to(torch.float32)

        return x

class HookVersion(nn.Module):
    def __init__(self, encoder):
        super(HookVersion, self).__init__()
        self.model = encoder
        self.layer_name = 'bn2' # changed this line
        self.activation = {}
        
        layer = dict([*self.model.visual.named_modules()])[self.layer_name]
        layer.register_forward_hook(self.get_activation(self.layer_name, self.activation))

    def get_activation(self, name, activations):
        def hook(model, input, output):
            activations[name] = output
        return hook

    def forward(self, x):
        _ = self.model.encode_image(x) 
        x = self.activation[self.layer_name]      
        x = torch.flatten(x).to(torch.float32)
        return x
    

fp_model = ForwardPassVersion(clip_model).to(device)
hook_model = HookVersion(clip_model).to(device)

fp_output = fp_model(image_input)
hook_output = hook_model(image_input)

print(torch.allclose(fp_output, hook_output, atol=1e-4))
</code></pre>
<p>it returns False. It seems the batchnorm layers here are acting differently. ReLU’s in later layers (such as layer3.0.relu3) also return different things (although they are equal before the first layer in this example). I get True for conv layers in late layers, however (such as layer3.0.conv3)</p>

ptrblck: <p>I see an inplace operation (<code>nn.ReLU(inplace=True)</code>) right after <code>bn2</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">clip_model.visual
Out[6]: 
ModifiedResNet(
  (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
  (bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu1): ReLU(inplace=True)
  (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu2): ReLU(inplace=True)
</code></pre>
<p>Also verified by checking the actual difference:</p>
<pre data-code-wrap="python"><code class="lang-python">(fp_output &lt; 0.).nonzero())
&gt; 172496

len((hook_output &lt; 0.).nonzero())
&gt; 0
</code></pre>