Load models give different results from original
❓ Questions and Help
What is your question?
What is the right way to retrieve a trained model and use it after load
class NER_Model(pl.LightningModule):
    def __init__(self, hyperparams, model_parameters, dataset_infos, extra_infos):
        super(NER_Model, self).__init__()

        # ---------- hyperparams
        self.learning_rate = hyperparams["learning_rate"]
        self.train_batch_size = hyperparams["train_batch_size"]
        self.eval_batch_size = hyperparams["eval_batch_size"]
        self.eps = hyperparams["eps"]
        self.seed = hyperparams["seed"]
        self.max_epochs = hyperparams["max_epochs"]
        self.MAX_LEN = hyperparams["MAX_LEN"]


        # ---------- model_parameters
        self.tokenizer_class = model_parameters["tokenizer_class"]
        self.tokenizer = self.tokenizer_class.tokenizer
        self.model_name = model_parameters["model_name"]
        self.output_attentions = model_parameters["output_attentions"]
        self.output_hidden_states = model_parameters["output_hidden_states"]
        
        # ---------- dataset_infos
        self.all_data = dataset_infos["all_data"]
        self.tags_infos_dict = dataset_infos["tags_infos_dict"]
        self.CustomDataset = dataset_infos["CustomDataset"]

        # ---------- extra_infos
        self.overfit = extra_infos["overfit"]
        self.sampler = extra_infos["sampler"]
        # ---------- other_infos
        self.predict_proba = torch.nn.Softmax(dim=1)
        self.step = "Experiment"
        
        # ---------- fixing seeds
        self.seed_everything()

        # # ---------- Model
        self.model = AutoModelForTokenClassification.from_pretrained(
          self.model_name,
          num_labels=len(self.tags_infos_dict["tag2idx"]),
          output_attentions = self.output_attentions,
          output_hidden_states = self.output_hidden_states
          )
        

        #Getting Dataset
        self.train_dataset = self.CustomDataset(
            self.all_data['X_train'], self.all_data['y_train'], self.tokenizer_class,self.tags_infos_dict,self.MAX_LEN,self.step 
        )
        self.valid_dataset = self.CustomDataset(
            self.all_data['X_valid'], self.all_data['y_valid'], self.tokenizer_class,self.tags_infos_dict,self.MAX_LEN,self.step 
        )
        self.test_dataset = self.CustomDataset(
            self.all_data['X_test'], self.all_data['y_test'], self.tokenizer_class,self.tags_infos_dict,self.MAX_LEN,self.step 
        )

        # ---------- Creating Dataframe for saving accuracy and loss
        self.df_performance_train_batch = pd.DataFrame(
            columns=["train_batch_loss", "train_batch_acc"]
        )
        self.df_performance_train_epoch = pd.DataFrame(
            columns=["train_epoch_loss", "train_epoch_acc"]
        )
        self.df_performance_valid_batch = pd.DataFrame(
            columns=["valid_batch_loss", "valid_batch_acc"]
        )
        self.df_performance_valid_epoch = pd.DataFrame(
            columns=["valid_epoch_loss", "valid_epoch_acc"]
        )
        


        # ---------- Creating comparision dataframe with expected and predicted results
        # self.response_columns = []
        # self.df_valid = pd.DataFrame(columns=self.response_columns)
        # self.df_test = pd.DataFrame(columns=self.response_columns)
        # self.valid_true_labels = []
        # self.valid_pred_labels = []
        self.test_true_tags = []
        self.test_pred_tags = []
        

    def predict(self, X: str):
        self.step = "Deployment"
        self.test_pred_tags = []

        batch = self.tokenizer.encode_plus(X, return_tensors="pt")
        batch["attention_masks"] = torch.ones_like(batch["input_ids"])
        batch = dict((key, input.to( self.device)) for key, input in batch.items())

        return self.test_step(batch, None)


    def retrieve_tags_on_prediction(self,predictions,inputs):
        predictions = predictions.detach().cpu().numpy()
        inputs =inputs.to('cpu').numpy()
        lables_list = []
        for input,prediction in zip(inputs,predictions):
          labels = []
          tokens = self.tokenizer.convert_ids_to_tokens(input)
          for token, pred_idx in zip(tokens, prediction):
              if not token.startswith("##"):
                  labels.append(self.tags_infos_dict['tag_values'][pred_idx])
          lables_list.append(labels)

        return lables_list
    

    def forward(self,input_ids,attention_mask=None,labels=None):
        if self.step == "Experiment":
            # loss, logits
            outputs = self.model(input_ids, 
                            token_type_ids=None,
                            attention_mask=attention_mask, 
                            labels=labels)
        if self.step == "Deployment":
                  # loss, logits
            outputs = self.model(input_ids)
  
          
        return outputs


    def training_step(self, batch, batch_nb):
        # batch
        inputs = batch['input_ids']
        mask = batch['attention_masks']
        targets = batch['tags']
                
        # fwd
        outputs = self.forward(inputs, mask, targets)
        loss = outputs[0]
        logits = outputs[1]

        # acc
        predictions = torch.argmax(logits, dim=2)
        pred_tags, true_tags = self.retrieve_tags(predictions,targets)
        acc = torch.tensor(accuracy_score(true_tags, pred_tags))

        # What to log
        tensorboard_logs = {"loss": loss, "acc": acc}

        self.df_performance_train_batch = self.df_performance_train_batch.append(
            pd.Series(
                [loss.item(), acc.item()], index=self.df_performance_train_batch.columns
            ),
            ignore_index=True,
        )

        return {
            "loss": loss,
            "train_acc_batch": acc,
            "train_loss_batch": loss,
            "log": tensorboard_logs,
        }

    def training_epoch_end(self, outputs):
        if not outputs:
            return {}

        temp_avg_loss_batch = [x["train_loss_batch"] for x in outputs]
        temp_avg_acc_batch = [x["train_acc_batch"] for x in outputs]

        avg_train_loss = torch.stack(temp_avg_loss_batch).mean()
        avg_train_acc = torch.stack(temp_avg_acc_batch).mean()

        self.df_performance_train_epoch = self.df_performance_train_epoch.append(
            pd.Series(
                [avg_train_loss.item(), avg_train_acc.item()],
                index=self.df_performance_train_epoch.columns,
            ),
            ignore_index=True,
        )

        tensorboard_logs = {
            "avg_train_acc": avg_train_acc,
            "avg_train_loss": avg_train_loss,
        }

        return {"avg_train_acc": avg_train_acc, "log": tensorboard_logs}

    def validation_step(self, batch, batch_nb):
        # batch
        inputs = batch['input_ids']
        mask = batch['attention_masks']
        targets = batch['tags']
                
        # fwd
        outputs = self.forward(inputs, mask, targets)
        loss = outputs[0]
        logits = outputs[1]

        # acc
        predictions = torch.argmax(logits, dim=2)
        pred_tags, true_tags = self.retrieve_tags(predictions,targets)
        # self.valid_true_labels.extend(true_tags)
        # self.valid_pred_labels.extend(pred_tags)
        acc = torch.tensor(accuracy_score(true_tags, pred_tags))

        self.df_performance_valid_batch = self.df_performance_valid_batch.append(
            pd.Series(
                [loss.item(), acc.item()], index=self.df_performance_valid_batch.columns
            ),
            ignore_index=True,
        )

        return {"valid_acc_batch": acc, "valid_loss_batch": loss}

    def validation_epoch_end(self, outputs):
        if not outputs:
            return {}
        temp_avg_loss_batch = [x["valid_loss_batch"] for x in outputs]
        temp_avg_acc_batch = [x["valid_acc_batch"] for x in outputs]

        avg_valid_loss = torch.stack(temp_avg_loss_batch).mean()
        avg_valid_acc = torch.stack(temp_avg_acc_batch).mean()

        self.df_performance_valid_epoch = self.df_performance_valid_epoch.append(
            pd.Series(
                [avg_valid_loss.item(), avg_valid_acc.item()],
                index=self.df_performance_valid_epoch.columns,
            ),
            ignore_index=True,
        )

        tensorboard_logs = {
            "avg_valid_acc": avg_valid_acc,
            "avg_valid_loss": avg_valid_loss,
        }

        return {"avg_valid_acc": avg_valid_acc, "log": tensorboard_logs}

    def test_step(self, batch, batch_nb):

        if self.step =="Experiment": 
          # batch
          inputs = batch['input_ids']
          mask = batch['attention_masks']
          targets = batch['tags']
                  
          # fwd
          outputs = self.forward(inputs, mask, targets)
          loss = outputs[0]
          logits = outputs[1]

          # acc
          predictions = torch.argmax(logits, dim=2)
          pred_tags, true_tags = self.retrieve_tags(predictions,targets)
          self.test_true_tags.extend(true_tags)
          self.test_pred_tags.extend(pred_tags)
          acc = torch.tensor(accuracy_score(true_tags, pred_tags))             
              
          final_return = {"test_acc_batch": acc}
        
        if self.step =="Deployment":
          # batch
          inputs = batch['input_ids']
          mask = batch['attention_masks']

          # fwd
          outputs = self.forward(inputs)
          logits = outputs[0]
          predictions = torch.argmax(logits, dim=2)

          pred_tags = self.retrieve_tags_on_prediction(predictions,inputs)
          final_return = pred_tags

        return final_return

    def test_epoch_end(self, outputs):
        if not outputs:
            return {}
          
        if self.step == "Experiment":
          avg_test_acc = torch.stack([x["test_acc_batch"] for x in outputs]).mean()

          tensorboard_logs = {"avg_test_acc": avg_test_acc}

          retorno = {"avg_test_acc": avg_test_acc, "log": tensorboard_logs}
        if self.step == "Deployment":
          retorno = outputs

        return retorno


    def configure_optimizers(self):
        optimizer = AdamW(
            [p for p in self.parameters() if p.requires_grad],
            lr = self.learning_rate,
            eps = self.eps
        )

        return optimizer
    
    def decode_tags(self,codes:np.ndarray):
      targets = []
      for elem in codes:
        targets.append(self.tags_infos_dict["idx2tag"][elem])
      return np.array(targets)
      
    def retrieve_tags(self,predictions,targets):
      targets = targets.detach().cpu().numpy()
      predictions = predictions.detach().cpu().numpy()
      pred_tags = [self.tags_infos_dict['tag_values'][p_i] for p, l in zip(predictions, targets) for p_i, l_i in zip(p, l) if self.tags_infos_dict['tag_values'][l_i] != self.tokenizer.pad_token]
      true_tags = [self.tags_infos_dict['tag_values'][l_i] for l in targets for l_i in l if self.tags_infos_dict['tag_values'][l_i] != self.tokenizer.pad_token]
      
      return pred_tags, true_tags

        
    def seed_everything(self):
        random.seed(self.seed)
        os.environ['PYTHONHASHSEED'] = str(self.seed)
        np.random.seed(self.seed)
        torch.manual_seed(self.seed)
        torch.cuda.manual_seed(self.seed)
        torch.cuda.manual_seed_all(self.seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    

    def gpu_mem_restore(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            except:
                type, val, tb = sys.exc_info()
                traceback.clear_frames(tb)
                raise type(val).with_traceback(tb) from None

        return wrapper              
            
    @gpu_mem_restore
    def train_dataloader(self):
        train_sampler = RandomSampler(self.train_dataset)
        train_dataloader = DataLoader(
                self.train_dataset,
                sampler=train_sampler,
                batch_size=self.train_batch_size,
                num_workers=cpu_count(),
                )
        return train_dataloader 

    @gpu_mem_restore
    def val_dataloader(self):
        valid_sampler = SequentialSampler(self.valid_dataset)
        val_dataloader = DataLoader(
                self.valid_dataset,
                sampler=valid_sampler,
                batch_size=self.eval_batch_size,
                num_workers=cpu_count(),
        )
        return val_dataloader
        
    @gpu_mem_restore
    def test_dataloader(self):
        test_sampler = SequentialSampler(self.test_dataset)
        test_dataloader = DataLoader(
                self.test_dataset,
                sampler=test_sampler,
                batch_size=self.eval_batch_size,
                num_workers=cpu_count(),
                )
        return test_dataloader

#### Training and testing a prediction (Wright Result!)
trainer = pl.Trainer(
    gpus=num_gpus,
    max_epochs=max_epochs,
    check_val_every_n_epoch=check_val_every_n_epoch,
    profiler=profiler,
    checkpoint_callback=checkpoint_callback,
    progress_bar_refresh_rate=progress_bar_refresh_rate,
    resume_from_checkpoint=resume_from_checkpoint,
    gradient_clip_val=gradient_clip_val,
)

model = NER_Model(
    hyperparams=hyperparams,
    model_parameters=model_parameters,
    dataset_infos=dataset_infos,
    extra_infos=extra_infos,
)

trainer.fit(model)
trainer.test(model)


#### Performing inference after training
model.predict("I love Brazilian beaches")
Result:
[['O',
  'O',
  'O',
  'O',
  'O',
  'Tag1',
  'Tag2,
  'O',
  'O']]

#### Saving **State Dict** and **Checkpoint**
model_path_dict = "model.pt"
model_path_checkpoint = "model.ckpt"
torch.save(model.state_dict(), model_path_dict)
trainer.save_checkpoint(model_path_checkpoint)


#### Testing lightning documentation sugestion (Wrong Result!)
device_cpu = torch.device('cpu') 
model = NER_Model.load_from_checkpoint(
    checkpoint_path = model_path_checkpoint,
    map_location=device,
    hyperparams=hyperparams,
    model_parameters=model_parameters,
    dataset_infos=dataset_infos,
    extra_infos=extra_infos,    
)
model.predict("I love Brazilian beaches")
Result:
[['Tag5',
  'Tag2',
  'O',
  'Tag4',
  'Tag4',
  'Tag3',
  'Tag1,
  'Tag5',
  'Tag4',]]



#### Testing loading the checkpoint (Wrong Result!)

model = NER_Model(
    hyperparams=hyperparams,
    model_parameters=model_parameters,
    dataset_infos=dataset_infos,
    extra_infos=extra_infos,    
)
storage = 'cpu'
checkpoint = torch.load(model_path_checkpoint, map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['state_dict'])

model.predict("I love Brazilian beaches")
Result:
[['Tag5',
  'Tag2',
  'O',
  'Tag4',
  'Tag4',
  'Tag3',
  'Tag1,
  'Tag5',
  'Tag4',]]


#### Testing loading directly the state dict (Wrong Result!)
model = NER_Model(
    hyperparams=hyperparams,
    model_parameters=model_parameters,
    dataset_infos=dataset_infos,
    extra_infos=extra_infos,    
)

device_cpu = torch.device('cpu')
model.load_state_dict(torch.load(model_path_dict,map_location=device_cpu)) 
model.eval()

model.predict("I love Brazilian beaches")
Result:
[['Tag5',
  'Tag2',
  'O',
  'Tag4',
  'Tag4',
  'Tag3',
  'Tag1,
  'Tag5',
  'Tag4,]]
Conclusion
Before saving and loading weights somewhere it works well (as expected).
After loading, in all of the tried strategies  it does not works as expected. The answer is wrong.
What's your environment?
Google Colab