Loading model from checkpoint results differ to loading model directly

sophiamaedler: <p>I am encountering issues where depending on how I load a model I obtain different results.</p>
<p>I have built a small test example which I have attached below that illustrates my problem. I have compared three different methods of loading the model:</p>
<ol>
<li>loading the model directly from hugging face</li>
<li>loading the model from a complete model checkpoint file</li>
<li>loading the model from a checkpoint file of the model state dict only</li>
</ol>
<p>I tested the models on the same randomly generated input images.</p>
<p>The weights and state dicts of the models using the 3 methods described are identical.<br>
The inference results generated by the 1st and the 3rd method are much more similar but still not the same.</p>
<p>Anyone have any ideas? I am completely at a loss what could be the result of this issue and how I should properly save and load my checkpoints to e.g. resume a training.</p>
<pre><code class="lang-auto">from transformers import ViTMAEForPreTraining
import torch
import gc
import pytorch_lightning as pl

import numpy as np
torch.cuda.empty_cache()

#define model 
class TransformerModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base")
        self.model.config.mask_ratio = 0.75
        
    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        imgs, labels = batch
        out = self(imgs)
        loss = out.loss
        self.log("train_loss", loss)
        return loss
    
    def validation_step(self, batch, batch_idx):
        imgs, labels = batch
        out = self(imgs)
        loss = out.loss
        self.log("val_loss", loss)
        return loss
        
    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr=0.001)
        return optimizer
    
    def unpatchify(self, x):
        p = self.model.vit.embeddings.patch_embeddings.patch_size[0]
        h = w = int(x.shape[1]**.5)
        
        assert h * w == x.shape[1]
        
        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
        x = torch.einsum('nhwpqc-&gt;nchpwq', x)
        imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))
        
        return imgs


#load model
model_1 = TransformerModel()
model_1.eval()

#save complete model as checkpoint
torch.save(model_1, "vit_mae_from_hugging_face.pth")

#save statedict only
torch.save(model_1.state_dict(), "vit_mae_from_hugging_face_statedict.pth")

#load model from checkpoint
model_2 = torch.load("vit_mae_from_hugging_face.cpkt")
model_2.eval()

#load state-dict only
model_3 = TransformerModel()
model_3.load_state_dict(torch.load("vit_mae_from_hugging_face_statedict.pth"))
model_3.eval()

#initialize hook to get intermediate outputs
layer_outputs = {}
def get_intermediate_1(module, input, output):
    layer_outputs['encoder_layernorm_1'] = output

def get_intermediate_2(module, input, output):
    layer_outputs['encoder_layernorm_2'] = output

def get_intermediate_3(module, input, output):
    layer_outputs['encoder_layernorm_3'] = output

model_1.model.vit.layernorm.register_forward_hook(get_intermediate_1)
model_2.model.vit.layernorm.register_forward_hook(get_intermediate_2)
model_3.model.vit.layernorm.register_forward_hook(get_intermediate_3)

#create a seeded random input tensor
torch.manual_seed(42)
tensor = torch.randn(8, 3, 224, 224)

#pass images to the two models
model_1(tensor)
model_2(tensor)
model_3(tensor)

outputs_1 = torch.mean(layer_outputs['encoder_layernorm_1'], axis = 1).detach().numpy()
outputs_2 = torch.mean(layer_outputs['encoder_layernorm_2'], axis = 1).detach().numpy()
outputs_3 = torch.mean(layer_outputs['encoder_layernorm_3'], axis = 1).detach().numpy()

#compare outputs
print("Output 1 vs Output2")
print((outputs_1 == outputs_2).all())
print("Output 2 vs Output3")
print((outputs_2 == outputs_3).all())
print("Output 1 vs Output3")
print((outputs_1 == outputs_3).all())

#compare model weights
print(((model_1.state_dict()['model.vit.embeddings.patch_embeddings.projection.weight']) == (model_2.state_dict()['model.vit.embeddings.patch_embeddings.projection.weight'])).all())
print(((model_1.state_dict()['model.vit.embeddings.patch_embeddings.projection.weight']) == (model_3.state_dict()['model.vit.embeddings.patch_embeddings.projection.weight'])).all())
print(((model_2.state_dict()['model.vit.embeddings.patch_embeddings.projection.weight']) == (model_3.state_dict()['model.vit.embeddings.patch_embeddings.projection.weight'])).all())

#compare complete state_dicts
for key in model_1.state_dict():
    value = (model_1.state_dict()[key] == model_2.state_dict()[key]).all()
    if not value:
        print(value)

for key in model_1.state_dict():
    pvalue = (model_1.state_dict()[key] == model_3.state_dict()[key]).all()
    if not value:
        print(value)
</code></pre>
<p>This returns:</p>
<pre><code class="lang-auto">Output 1 vs Output2
False
Output 2 vs Output3
False
Output 1 vs Output3
False
tensor(True)
tensor(True)
tensor(True)
</code></pre>
<p>Numerically the outputs generated by approach 1 and approach 3 are much more similar than those generated by approach 2.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/f/5f8b63b57c5fb949f5ab5ca1980be73ad1741806.png" data-download-href="https://discuss.pytorch.org/uploads/default/5f8b63b57c5fb949f5ab5ca1980be73ad1741806" title="Screenshot 2024-11-04 at 18.41.56"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/f/5f8b63b57c5fb949f5ab5ca1980be73ad1741806_2_581x500.png" alt="Screenshot 2024-11-04 at 18.41.56" data-base62-sha1="dDe017Vh3hxUD5lCZFWIOZjiLAi" width="581" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/f/5f8b63b57c5fb949f5ab5ca1980be73ad1741806_2_581x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/f/5f8b63b57c5fb949f5ab5ca1980be73ad1741806_2_871x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/f/5f8b63b57c5fb949f5ab5ca1980be73ad1741806_2_1162x1000.png 2x" data-dominant-color="EDEDED"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-11-04 at 18.41.56</span><span class="informations">1282×1102 84.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <aside class="quote no-group" data-username="sophiamaedler" data-post="1" data-topic="212514">
<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/sophiamaedler/48/66166_2.png" class="avatar"> sophiamaedler:</div>
<blockquote>
<p>loading the model from a complete model checkpoint file</p>
</blockquote>
</aside>
<p>Serializing the entire model is not recommended as it can easily break (e.g. if the source files were changed). I don’t know why this approach causes the large numerical mismatch in your example, as I cannot download the checkpoint (it’s terribly slow).</p>

sophiamaedler: <p>Thanks for your insights. We managed to find the source of the issue after a helpful tip from a colleague I shared the post with.</p>
<p>The VitMAE model we are using here randomly masks part of the image which it tries to reconstruct. The <code>mask_ratio</code> determines which fraction of the image is masked which was the same across models, but the actual parts of the image that are masked is random. So even on identical images for <code>mask_ratio != 0</code> the returned results are not the same.</p>
<p>By setting the mask_ratio to 0 (and thus removing any remaining randomness) we could recreate our results down to the 1e-8 range which is probably a result of calculation inaccuracies on the GPU resulting from optimisation algorithms. So at the end of the day it was a user error on our part.</p>
<p>What we did notice and maybe you can provide some insights on is even if I set the mask ratio to 0 before saving the checkpoint when I loaded the model directly from the checkpoint using serialisation of the entire model the mask_ratio was back to 0.75 after loading which I manually had to set to 0 again.</p>
<p>For anyone interested here the code with the mask_ratios set that leads to reproducible results:</p>
<pre><code class="lang-auto">from transformers import ViTMAEForPreTraining
import torch
import gc
import pytorch_lightning as pl

import numpy as np
torch.cuda.empty_cache()

#define model 
class TransformerModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base")
        self.model.config.mask_ratio = 0.0
        
    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        imgs, labels = batch
        out = self(imgs)
        loss = out.loss
        self.log("train_loss", loss)
        return loss
    
    def validation_step(self, batch, batch_idx):
        imgs, labels = batch
        out = self(imgs)
        loss = out.loss
        self.log("val_loss", loss)
        return loss
        
    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr=0.001)
        return optimizer
    
    def unpatchify(self, x):
        p = self.model.vit.embeddings.patch_embeddings.patch_size[0]
        h = w = int(x.shape[1]**.5)
        
        assert h * w == x.shape[1]
        
        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
        x = torch.einsum('nhwpqc-&gt;nchpwq', x)
        imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))
        
        return imgs


#load model
model_1 = TransformerModel()
model_1.eval()

#save complete model as checkpoint
torch.save(model_1, "vit_mae_from_hugging_face.pth")

#save statedict only
torch.save(model_1.state_dict(), "vit_mae_from_hugging_face_statedict.pth")

#load model from checkpoint
model_2 = torch.load("vit_mae_from_hugging_face.cpkt")
print(model_2.model.config.mask_ratio)
model_2.model.config.mask_ratio = 0
model_2.eval()

#load state-dict only
model_3 = TransformerModel()
model_3.load_state_dict(torch.load("vit_mae_from_hugging_face_statedict.pth"))
model_3.eval()

#initialize hook to get intermediate outputs
layer_outputs = {}
def get_intermediate_1(module, input, output):
    layer_outputs['encoder_layernorm_1'] = output

def get_intermediate_2(module, input, output):
    layer_outputs['encoder_layernorm_2'] = output

def get_intermediate_3(module, input, output):
    layer_outputs['encoder_layernorm_3'] = output

model_1.model.vit.layernorm.register_forward_hook(get_intermediate_1)
model_2.model.vit.layernorm.register_forward_hook(get_intermediate_2)
model_3.model.vit.layernorm.register_forward_hook(get_intermediate_3)

#create a seeded random input tensor
torch.manual_seed(42)
tensor = torch.randn(8, 3, 224, 224)

#pass images to the two models
model_1(tensor)
model_2(tensor)
model_3(tensor)

outputs_1 = torch.mean(layer_outputs['encoder_layernorm_1'], axis = 1).detach().numpy()
outputs_2 = torch.mean(layer_outputs['encoder_layernorm_2'], axis = 1).detach().numpy()
outputs_3 = torch.mean(layer_outputs['encoder_layernorm_3'], axis = 1).detach().numpy()

#compare outputs
print("Output 1 vs Output2")
print((outputs_1 == outputs_2).all())
print("Output 2 vs Output3")
print((outputs_2 == outputs_3).all())
print("Output 1 vs Output3")
print((outputs_1 == outputs_3).all())

#compare model weights
print(((model_1.state_dict()['model.vit.embeddings.patch_embeddings.projection.weight']) == (model_2.state_dict()['model.vit.embeddings.patch_embeddings.projection.weight'])).all())
print(((model_1.state_dict()['model.vit.embeddings.patch_embeddings.projection.weight']) == (model_3.state_dict()['model.vit.embeddings.patch_embeddings.projection.weight'])).all())
print(((model_2.state_dict()['model.vit.embeddings.patch_embeddings.projection.weight']) == (model_3.state_dict()['model.vit.embeddings.patch_embeddings.projection.weight'])).all())

#compare complete state_dicts
for key in model_1.state_dict():
    value = (model_1.state_dict()[key] == model_2.state_dict()[key]).all()
    if not value:
        print(value)

for key in model_1.state_dict():
    pvalue = (model_1.state_dict()[key] == model_3.state_dict()[key]).all()
    if not value:
        print(value)
</code></pre>
<p>Which returns:</p>
<pre><code class="lang-auto">0.75
Output 1 vs Output2
False
Output 2 vs Output3
False
Output 1 vs Output3
False
tensor(True)
tensor(True)
tensor(True)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/d/7d24e321ae9cec42ce8c26b0fcfce6947d859231.png" data-download-href="https://discuss.pytorch.org/uploads/default/7d24e321ae9cec42ce8c26b0fcfce6947d859231" title="Screenshot 2024-11-06 at 12.09.32"><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/d/7d24e321ae9cec42ce8c26b0fcfce6947d859231.png" alt="Screenshot 2024-11-06 at 12.09.32" data-base62-sha1="hR4IhUFTFBKwtigeL35xkt0Kn1D" width="575" height="500" data-dominant-color="ECECEC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-11-06 at 12.09.32</span><span class="informations">632×549 23.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>