Encountering an Issue While Building a Queryable Protein Sequence Database

[pengjk689]

Dear,

When I attempt to build the tmvec database using the entire human protein data from the UniProt database (208,022 entries), I encounter the following issue:

Traceback (most recent call last): File "/home/pengjiak/miniconda3/envs/tmvec/bin/tmvec-build-database", line 110, in encoded_database = encode(flat_seqs, model_deep, model, tokenizer, device) File "/home/pengjiak/miniconda3/envs/tmvec/lib/python3.9/site-packages/tm_vec/tm_vec_utils.py", line 61, in encode protrans_sequence = featurize_prottrans(sequences[i:i+1], model, tokenizer, device) File "/home/pengjiak/miniconda3/envs/tmvec/lib/python3.9/site-packages/tm_vec/tm_vec_utils.py", line 24, in featurize_prottrans embedding = model(input_ids=input_ids, attention_mask=attention_mask) File "/home/pengjiak/miniconda3/envs/tmvec/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl return forward_call(*args, kwargs) File "/home/pengjiak/miniconda3/envs/tmvec/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py", line 1964, in forward encoder_outputs = self.encoder( File "/home/pengjiak/miniconda3/envs/tmvec/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl return forward_call(*args, *kwargs) File "/home/pengjiak/miniconda3/envs/tmvec/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py", line 1123, in forward layer_outputs = layer_module( File "/home/pengjiak/miniconda3/envs/tmvec/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl return forward_call(args, kwargs) File "/home/pengjiak/miniconda3/envs/tmvec/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py", line 695, in forward self_attention_outputs = self.layer[0]( File "/home/pengjiak/miniconda3/envs/tmvec/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl return forward_call(*args, *kwargs) File "/home/pengjiak/miniconda3/envs/tmvec/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py", line 602, in forward attention_output = self.SelfAttention( File "/home/pengjiak/miniconda3/envs/tmvec/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl return forward_call(args, **kwargs) File "/home/pengjiak/miniconda3/envs/tmvec/lib/python3.9/site-packages/transformers/models/t5/modeling_t5.py", line 552, in forward position_bias = position_bias + mask # (batch_size, n_heads, seq_length, key_length) torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 25.09 GiB (GPU 0; 79.18 GiB total capacity; 55.53 GiB already allocated; 22.96 GiB free; 55.73 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

It's worth noting that I encounter this issue even when attempting to build the database with 20,000 protein FASTA files, which is quite perplexing. Looking forward to your response.

[pengjk689]

here my code: tmvec-build-database \ --input-fasta human_swissProt.fa \ --tm-vec-model ${path}/tm_vec_cath_model.ckpt \ --tm-vec-config-path ${path}/tm_vec_cath_model_params.json \ --output human_swissProt_database \ --protrans-model ${path}/prot_t5_xl_uniref50 \ --device 'gpu' \

[mortonjt]

Hi, we cannot currently do that from the CLI -- you'll need to batch it into smaller chunks, encode and create the database from the encodings (otherwise you'll cram too much into gpu memory). We'll try to streamline this in a follow up release