Model(x) not matching head(backbone(x))

aniketmaurya: <p>Hi, I have a model which is composed of the backbone (Resnet50) + head(linear layer).<br>
During training, I have frozen backbone so my understanding is that <code>model(x) = head(backbone(x))</code> but it is not!</p>
<p>Before testing, I am saving head and backbone separately and then load as-</p>
<pre><code class="lang-auto">head = torch.load('head.pt')
backbone = torch.load('backbone.pt')

model = torch.load('model.pt')

assert model(x) == head(backbone(x))
</code></pre>
<p>The prediction final results are correct mostly but there is a significant change in raw logit values.</p>
<p>EDIT:<br>
all three- model, backbone, head were set to eval before testing and <code>torch.no_grad</code> was used.</p>

aniketmaurya: <p>this is my forward definition -</p>
<pre><code class="lang-auto">class Net(nn.Module):

  def __init__(self):
      self.backbone = timm.create_model('resnet50', num_classes=0).eval()
      self.head = nn.Linear(2048, num_classes)

  def forward(self, x):
      with torch.no_grad():
         x = self.backbone(x)
      x = self.head(x)
   return x
</code></pre>

aniketmaurya: <p>I guess backbone parameters are being updated somehow even when it’s in <code>eval</code> model and <code>torch.no_grad</code> is used.</p>
<p>Because this is failing -</p>
<pre><code class="lang-python">timm_backbone = timm.create_model('resnet50', num_classes=0).eval()

assert model.backbone(x) == timm_backbone(x)
</code></pre>

ptrblck: <p>Directly comparing floating point values might fail due to the limited floating point precision, which would be visible if e.g. a non-deterministic algorithm is used so use <code>torch.allclose</code> instead.</p>
<aside class="quote no-group" data-username="aniketmaurya" data-post="1" data-topic="131412">
<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/aniketmaurya/48/43800_2.png" class="avatar"> aniketmaurya:</div>
<blockquote>
<p>The prediction final results are correct mostly but there is a significant change in raw logit values.</p>
</blockquote>
</aside>
<p>How large is this significant change if you <code>print</code> the <code>abs().max()</code> of the difference?</p>

aniketmaurya: <p>Mostly they have a difference of around 0.2</p>

ptrblck: <p>Thanks for the update!<br>
I don’t know how your models are defined and if you are missing functional API calls or if there is any other issue, but I cannot reproduce the issue using the <code>torchvision</code> ResNet152 by comparing the entire model execution with a per-layer one:</p>
<pre><code class="lang-python">model = models.resnet152()
model.eval()

x = torch.randn(1, 3, 224, 224)

# entire model execution
out1 = model(x)

# per-layer execution
x = model.conv1(x)
x = model.bn1(x)
x = model.relu(x)
x = model.maxpool(x)
x = model.layer1(x)
x = model.layer2(x)
x = model.layer3(x)
x = model.layer4(x)
x = model.avgpool(x)
x = torch.flatten(x, 1)
out2 = model.fc(x)

# compare
print((out1 - out2).abs().max())
&gt; tensor(0., grad_fn=&lt;MaxBackward1&gt;)
</code></pre>

aniketmaurya: <p>thank you for your quick response <img src="https://discuss.pytorch.org/images/emoji/apple/pray/2.png?v=10" title=":pray:t2:" class="emoji" alt=":pray:t2:">.<br>
I think the issue is coming when I train the model somehow backbone is getting updated even though it is in eval mode.</p>
<pre><code class="lang-python">class Net(nn.Module):

  def __init__(self):
      self.backbone = timm.create_model('resnet50', num_classes=0).eval()
      self.head = nn.Linear(2048, num_classes)

  def forward(self, x):
      with torch.no_grad():
         x = self.backbone(x)
      x = self.head(x)
   return x

def train(net, dataloader):
    """model trained with crossentropy here"""
  
</code></pre>
<p>After training, what I want is to only save the <code>head</code>. I will do this with a couple of tasks and use a single pretrained frozen backbone for extracting features and use the head of each task for getting the probabilities.</p>

ptrblck: <p>I’m not sure how the training fits into this use case, since you are comparing the outputs directly in your first post. However, if you think the backbone is updated, check the <code>.grad</code> attributes of the parameters of <code>model.backbone</code> after the <code>backward</code> call and verify that they are set to <code>None</code>.</p>

aniketmaurya: <p>I verified the below two things -</p>
<pre data-code-wrap="python"><code class="lang-python">for p in model.backbone.parameters():
    assert p.grad==None


forzen_backbone = timm.create_model('ssl_resnet50', True, num_classes=0).eval()
# model.backbone was also frozen during training
for p, p2 in zip(model.backbone.parameters(), forzen_backbone.parameters()):
    print(torch.all(torch.isclose(p, p2)))
</code></pre>
<blockquote>
<p>Prints all True</p>
</blockquote>

aniketmaurya: <p>FYI, that Batch Normalisation tracking was not disabled on model.eval() or with torch.no_grad. After setting BN tracking to False the issue was gone.</p>