I am trying to reload a fine-tuned DistilBertForTokenClassification model. After using the Trainer to train the downloaded model, I save the model with trainer.save_model(),
## Model training
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased", use_fast=True)
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=512)
trainer = Trainer(
...
)
trainer.train()
trained_model_save_path = 'some_path/on_disk'
torch.save(model.state_dict(), trained_model_save_path)
sample_test_output = get_prediction(model, sample_test_text= "sample_test_text", MAX_LEN)
In the inference (a separate ipython script), I tried to load this model from disk. For the same sample_test_text, sample_test_output is inconsistent and worse than that in the model training. So I suspect the model weights were loaded incorrectly in inference. In the model loading part, I tried to mimic the hugginface tutorial: model = TheModelClass(*args, **kwargs)
## Model Inference
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=512)
model.load_state_dict(torch.load(trained_model_save_path))
model.eval()
sample_test_output = get_prediction(model, sample_test_text= "sample_test_text", MAX_LEN)
Seems there is another question reported similar issue but no effective answers yet: Saving and reload huggingface fine-tuned transformer
