theislab / chemCPA

Code for "Predicting Cellular Responses to Novel Drug Perturbations at a Single-Cell Resolution", NeurIPS 2022.
https://arxiv.org/abs/2204.13545
MIT License
99 stars 24 forks source link

Architecture surgery for transfer experiements #62

Closed MxMstrmn closed 2 years ago

MxMstrmn commented 2 years ago

We need to add the code to support architecture surgery. This covers scenarios when we have a pretrained model on X genes but only a subset of X in the new dataset on which we want to transfer. More specifically, this refers to a scenario where we only use say 500 from the 974 lincs genes during the transfer task.

siboehm commented 2 years ago

(I'll keep updating this comment) This requires the following changes:

  1. Store the gene_id inside the model.pt-Checkpoints. This is necessary st we know which genes & gene orderings the model was trained on.
  2. Write the code that performs the surgery:
    1. First, compare the gene_ids of the finetuning dataset with the (stored) gene ids of the dataset that we pretrained on
    2. For genes that are missing in the finetuning dataset: Delete the respective rows (Encoder) and Columns (Decoder) from the state_dict, then load the weights.
    3. For genes that are new in the finetuning dataset: Add randomly initialized rows (Encoder) and Columns (Decoder) to the state_dict, then load the weights.
  3. Run the finetuning with some reasonable hyperparameters
siboehm commented 2 years ago

For now we'll just do gene set expansion, which we can cover well enough by just adding another layer. Therefore we've decided to close this issue in favor of working on #71