In silico Perturbation Out of Memory
I realise this topic has been discussed before, but I have gone through previous responses and none of those solutions worked for me. I would like to set max_ncells to 'None' but I can only go as far as 6500. I tried reducing cell_inds_to_perturb and forward_batch_size, but neither was sufficient. The OOM occurs before saving disk so clearing memory more often has no effect. I am using an AWS instance with 24GB GPU RAM. Thanks in advance!
isp = InSilicoPerturber(
perturb_type="delete",
perturb_rank_shift=None,
genes_to_perturb=genes,
combos=0,
anchor_gene=None,
model_type="CellClassifier",
num_classes=3,
emb_mode="cell",
cell_emb_style="mean_pool",
filter_data={"cell_type":["Cardiomyocyte1","Cardiomyocyte2","Cardiomyocyte3"]},
cell_states_to_model={
'state_key': 'disease',
'start_state': 'dcm',
'goal_state': 'nf',
'alt_states': ['hcm']
},
cell_inds_to_perturb=[{'start': 1, 'end': 1000}],
max_ncells=6500,
emb_layer=0,
forward_batch_size=3,
nproc=16
)
isp.perturb_data(
"./CellClassifier_cardiomyopathies/",
"./human_dcm_hcm_nf.dataset",
"./perturb_out/",
"prefix_"
)
Error:
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /home/ubuntu/Geneformer/examples/in_silico_perturbation.py:44 in <module> │
│ │
│ 41 │
│ 42 cell_ind_ranges_list=[{'start': 1, 'end': 1000}] │
│ 43 for x in cell_ind_ranges_list: │
│ ❱ 44 │ perform_perturb(x) │
│ 45 │
│ 46 #with ProcessPoolExecutor(max_workers=len(cell_ind_ranges_list)) as executor: │
│ 47 # executor.map(perform_perturb, cell_ind_ranges_list) │
│ │
│ /home/ubuntu/Geneformer/examples/in_silico_perturbation.py:34 in perform_perturb │
│ │
│ 31 │ ) │
│ 32 │ │
│ 33 │ # outputs intermediate files from in silico perturbation │
│ ❱ 34 │ isp.perturb_data( │
│ 35 │ │ "./CellClassifier_cardiomyopathies/", │
│ 36 │ │ "/home/ubuntu/Geneformer/example_input_files/cell_classification/disease_classif │
│ 37 │ │ "./perturb_out/", │
│ │
│ /home/ubuntu/miniconda3/lib/python3.10/site-packages/geneformer/in_silico_perturber.py:958 in │
│ perturb_data │
│ │
│ 955 │ │ │ │ │ raise │
│ 956 │ │ │ # get dictionary of average cell state embeddings for comparison │
│ 957 │ │ │ downsampled_data = downsample_and_sort(filtered_input_data, self.max_ncells) │
│ ❱ 958 │ │ │ state_embs_dict = get_cell_state_avg_embs(model, │
│ 959 │ │ │ │ │ │ │ │ │ │ │ │ │ downsampled_data, │
│ 960 │ │ │ │ │ │ │ │ │ │ │ │ │ self.cell_states_to_model, │
│ 961 │ │ │ │ │ │ │ │ │ │ │ │ │ layer_to_quant, │
│ │
│ /home/ubuntu/miniconda3/lib/python3.10/site-packages/geneformer/in_silico_perturber.py:303 in │
│ get_cell_state_avg_embs │
│ │
│ 300 │ │ │ torch.cuda.empty_cache() │
│ 301 │ │ │
│ 302 │ │ state_embs = torch.cat(state_embs_list) │
│ ❱ 303 │ │ avg_state_emb = mean_nonpadding_embs(state_embs, torch.Tensor(original_lens).to( │
│ 304 │ │ avg_state_emb = torch.mean(avg_state_emb, dim=0, keepdim=True) │
│ 305 │ │ state_embs_dict[possible_state] = avg_state_emb │
│ 306 │ return state_embs_dict │
│ │
│ /home/ubuntu/miniconda3/lib/python3.10/site-packages/geneformer/in_silico_perturber.py:587 in │
│ mean_nonpadding_embs │
│ │
│ 584 │ mask = mask.unsqueeze(2).expand_as(embs) │
│ 585 │ │
│ 586 │ # use the mask to zero out the embeddings in padded areas │
│ ❱ 587 │ masked_embs = embs * mask.float() │
│ 588 │ │
│ 589 │ # sum and divide by the lengths to get the mean of non-padding embs │
│ 590 │ mean_embs = masked_embs.sum(1) / original_lens.view(-1, 1).float() │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
Thank you for your interest in Geneformer and for your question! In this case the OOM is from calculating the target embedding positions for the cell states being modeled. When setting max_ncells to None, all of the cells in the dataset are used to calculate these initial positions even though the cell_inds_to_perturb subsets the cells used for the perturbation tests. We will look into this memory error to see if it’s possible to obtain these initial embedding positions in a more memory preserving manner, but in the meantime, you could obtain the goal embedding positions separately and modify the code to supply them to each of your runs with subsets of cells. At the end, you can put all the results into the same directory for the stats module to analyze them together.
As a quick fix:
downsampled_data = downsample_and_sort(filtered_input_data, self.max_ncells)
num_shards = 40 # data will be split into 40
shard_num= 0 # which shard will be analysed
shards = []
for i in range(num_shards):
shard = downsampled_data.shard(num_shards=num_shards, index=i)
shards.append(shard)
print(f'shard num is {shard_num}')
downsampled_data = shards[shard_num]
Here I'm splitting the data into 40 'shards' and will analyse them separately. With my current set up (g4dn.metal AWS) this will take an estimated 24 days though, so would greatly appreciate any pointers to make this more efficient!
Edit: Running the entire dataset seems to be a little overkill. FDR and cosine shift values dont change much after ~20% of the cells are analysed, which reduces the analysis time significantly.