A MTLClassifier.load_and_evaluate_test_model error
Dear Geneformer Team,
I am trying to fine-tune a Geneformer model pretrained on mouse data (mousegeneformer). I used the MTLClassifier and was able to successfully run:
MTLClassifier.run_manual_tuning()
However, I cannot run:
MTLClassifier.load_and_evaluate_test_model()
It fails with the following error:
RuntimeError: Error(s) in loading state_dict for GeneformerMultiTask:
Unexpected key(s) in state_dict: "attention_pool.attention_weights".
My Geneformer installation is from the official repository on Hugging Face. I only modified the TranscriptomeTokenizer class to allow the V1 model to work without some default parameters:（because mousegeneformer is a V1 model）
elif self.model_version == "V1":
self.model_input_size = 2048
self.special_token = False
# from . import ENSEMBL_MAPPING_FILE_30M, GENE_MEDIAN_FILE_30M, TOKEN_DICTIONARY_FILE_30M
gene_median_file = gene_median_file #it used to be ENSEMBL_MAPPING_FILE_30M
token_dictionary_file = token_dictionary_file #......
gene_mapping_file = gene_mapping_file #......
My MTLClassifier was initialized as follows:
task_columns = ["predicted_cell_type", "treatment"]
hyperparameters = {
"learning_rate": 1e-4,
"warmup_ratio": 0.15,
"weight_decay": 0.05,
"dropout_rate": 0.1,
"lr_scheduler_type": "cosine",
"task_weights": [1,1],
"max_layers_to_freeze": 2
}
mc = MTLClassifier(
task_columns=task_columns,
study_name="MTLClassifier_example",
pretrained_path=pretrained_path,
train_path=train_path,
val_path=val_path,
test_path=test_path,
model_save_path=model_save_path,
results_dir=results_dir,
tensorboard_log_dir=tensorboard_log_dir,
use_manual_hyperparameters=True,
manual_hyperparameters=hyperparameters,
epochs=1, # Suggested to prevent overfitting
batch_size=32, # Adjust based on GPU memory
seed=42
)
The traceback is as follows:
mc.load_and_evaluate_test_model()
No unique_cell_id column in train dataset
Collected 0 cell IDs from training dataset
No unique_cell_id column in validation dataset
Collected 37013 cell IDs from validation dataset
No unique_cell_id column in test dataset
Collected 37037 cell IDs from test dataset
Loaded hyperparameters:
learning_rate: 0.0001
warmup_ratio: 0.15
weight_decay: 0.05
dropout_rate: 0.1
lr_scheduler_type: cosine
normalized_task_weights: [1, 1]
max_layers_to_freeze: 2
use_task_weights: True
use_attention_pooling: True
RuntimeError Traceback (most recent call last)
...
RuntimeError: Error(s) in loading state_dict for GeneformerMultiTask:
Unexpected key(s) in state_dict: "attention_pool.attention_weights".
Could you please advise how I can resolve this issue so that load_and_evaluate_test_model() works with my fine-tuned mouse model?
Thank you very much for your help.
Best regards
Thank you for your question. It's possible the mouse model had use_attention_pooling set to False during training, so you may try setting this to False when you load and evaluate the saved model. With regard to using a V1 model, we just would recommend ensuring that all downstream work also is set for the expectation of the V1 model with regard to the token dictionary, input size, etc.