Why my model behaves differently at each load?

NouRed: <h3>
<a name="tldr-1" class="anchor" href="#tldr-1"></a>TL;DR</h3>
<p>The fine-tuned BertModel on NER task behaves differently at each load of the <code>.bin</code> file.</p>
<h3>
<a name="example-2" class="anchor" href="#example-2"></a>Example:</h3>
<p>When the model finishes training:</p>
<ul>
<li>Input<br>
<code>I am John and I work at Hugging-Face</code>
</li>
<li>output<br>
<code>[(I, O), (am, O), (John, PER), (and, O), (I, O), (work, O), (at, O), (Hugging-Face, ORG)]</code>
</li>
</ul>
<p>After stopping the notebook session and loading the model:</p>
<ul>
<li>Input<br>
<code>I am John and I work at Hugging-Face</code>
</li>
<li>output<br>
<code>[(I, PER), (am, PER), (John, PER), (and, PER), (I, ORG), (work, PER), (at, O), (Hugging-Face, PER)]</code>
</li>
</ul>
<h4>
<a name="environment-3" class="anchor" href="#environment-3"></a><strong>Environment:</strong>
</h4>
<ul>
<li><strong>Colab Pro +</strong></li>
<li><strong>Transformers == 4.23.1</strong></li>
<li><strong>Torch == 1.12.1</strong></li>
</ul>
<h4>
<a name="description-4" class="anchor" href="#description-4"></a>Description</h4>
<p>I am currently facing an issue with my NER model based on <code>BertModel</code> from the Transformers library and inspired from the <code>BertForTokenClassification</code> code base.</p>
<p>Indeed, the issue is the following, after training and evaluating my model I end up with a well-performing model with a validation accuracy greater than 96%. The problem is that when I save the model and load it for inference it gives different results, yes it gives different predictions (bad) each time it is loaded. It should be noted that when the model has finished training the predictions are good, but when I stop the notebook session and start another one and then load my best model saved, it behaves differently.</p>
<p><strong>Model Architecture</strong>:</p>
<pre><code class="lang-python">class NerBertModel(nn.Module):
  def __init__(self, id2label, label2id, num_labels):
    super(PhenoBertModel, self).__init__()
    self.id2label = id2label
    self.label2id = label2id
    self.num_labels = num_labels
    
    self.bert = Config.MODEL
    classifier_dropout = (
            Config.CONFIG.classifier_dropout if Config.CONFIG.classifier_dropout is not None else Config.CONFIG.hidden_dropout_prob
        )
    self.dropout = nn.Dropout(classifier_dropout)

    self.classifier = nn.Linear(Config.CONFIG.hidden_size, num_labels)

  def forward(self, 
              input_ids: Optional[torch.Tensor] = None, 
              attention_mask: Optional[torch.Tensor] = None, 
              token_type_ids: Optional[torch.Tensor] = None,
              labels: Optional[torch.Tensor] = None):
    
    outputs = self.bert(input_ids, attention_mask)

    sequence_output = outputs[0]
    
    sequence_output = self.dropout(sequence_output)
    logits = self.classifier(sequence_output)

    loss = None
    if labels is not None:
      loss_fct = nn.CrossEntropyLoss()
      loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
    
    return loss, logits
</code></pre>
<p><strong>Saved the model using:</strong></p>
<pre><code class="lang-python">torch.save(model.state_dict(), Config.MODEL_PATH)
</code></pre>
<p><strong>Loaded the model using:</strong></p>
<pre><code class="lang-python">model = NerBertModel(id2label, label2id, num_labels=len(id2label))

model.load_state_dict(
 torch.load(
  Config.MODEL_PATH, # model.bin file
  map_location=torch.device(Config.DEVICE) 
)

model.to(Config.DEVICE)
model.eval()
</code></pre>
<p>The same problem occurs also when using the standard NER model <code>BertForTokenClassfication</code> from the Transformers library directly while saving it and loading as follows:</p>
<pre><code class="lang-python"># Save best model
mode.save_pretrained("model_path")

# Load the model
BertForTokenClassfication.from_pretrained("model_path").
</code></pre>
<p>The <strong>seed</strong> function I am using</p>
<pre><code class="lang-python">def seed_torch(seed=42):
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
seed_torch()
</code></pre>

NouRed: <h3>
<a name="problem-fixed-1" class="anchor" href="#problem-fixed-1"></a>Problem Fixed</h3>
<p>I fixed it. The problem is that each time I run the notebook, <code>unique_labels</code> contains the labels in a different order compared to the previous notebook session, so I end up with different encoding of the labels and this is due to the <code>set()</code> method, which I used to get unique labels and then encode them dynamically as shown in the snippet of code below:</p>
<pre><code class="lang-python">unique_labels = set([label for label in data["token_labels"].values for label in labels])
label2id = {k: v for v, k in enumerate(unique_labels)}
id2label = {v: k for v, k in enumerate(unique_labels)}
</code></pre>
<p>It should be noted that the problem occurred even when:</p>
<ol>
<li>Using the standard <code>BertForTokenClassification</code> model from the Hugging-Face transformers library while using <code>save_pretrained()</code> and <code>from_pretrained()</code>. However, It is recommended to save and load the best model using <code>save_pretrained()</code> and <code>from_pretrained()</code> respectively when it comes to a model based on the Hugging-Face transformers library.</li>
<li>Running the notebook on the local host.</li>
</ol>
<p>So just try to avoid using <code>set()</code> or <strong>sort its output before label encoding</strong>, thus you always end up with the same label encoding.</p>