Inquiry about GPU Memory Requirements for Running Pre-trained Model

[OrientCCC]

Hello @Layne-Huang,

First of all, thank you for developing and sharing such an impactful project！

I am interested in leveraging the pre-trained model you provided for sample novel molecules given seed fragments using the following command as mentioned in the documentation:

python -u sample_frag.py --ckpt 500.pt --pdb_path ./mol_data/receptor.pdb --mol_file ./mol_data/receptor_ligand.sdf --keep_index 1 2 3 4 5 6 7 --num_atom 16 --num_samples 50 --sampling_type generalized

However, I encountered a CUDA out of memory error:

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 243.00 GiB. GPU 0 has a total capacty of 10.75 GiB of which 8.49 GiB is free. Including non-PyTorch memory, this process has 2.26 GiB memory in use. Of the allocated memory 503.59 MiB is allocated by PyTorch, and 1.58 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

It seems that the model requires a significant amount of GPU memory, and I'm curious about the typical GPU memory requirements for running this pre-trained model with the provided command in common scenarios.

Could you please provide some guidance on the expected GPU memory usage for running the pre-trained model with the provided command? Any suggestions or recommendations for managing GPU memory effectively would also be greatly appreciated.

Thank you very much for your time and assistance!

[FeilongWuHaa]

Maybe you can decrease the batch_size

python -u sample_frag.py --ckpt 500.pt --pdb_path ./mol_data/receptor.pdb --mol_file ./mol_data/receptor_ligand.sdf --keep_index 1 2 3 4 5 6 7 --num_atom 16 --num_samples 50 --sampling_type generalized --batch_size 16

[OrientCCC]

Thank you @FeilongWuHaa for the suggestion to decrease the batch_size. I've tried adjusting it to a smaller size, and the model works during testing.

This workaround is helpful for now, but for planning and scaling purposes, I'm still interested in understanding the typical GPU memory requirements for this model under common usage scenarios.

It would be great to hear the author's insights or anyone else's experiences regarding the GPU usage for this model, especially when dealing with larger batch sizes or more complex inputs. This information would be valuable for configuring computation resources effectively.

Thanks again for your helpful advice!

[Layne-Huang]

Thank you @FeilongWuHaa for the suggestion to decrease the batch_size. I've tried adjusting it to a smaller size, and the model works during testing.

This workaround is helpful for now, but for planning and scaling purposes, I'm still interested in understanding the typical GPU memory requirements for this model under common usage scenarios.

It would be great to hear the author's insights or anyone else's experiences regarding the GPU usage for this model, especially when dealing with larger batch sizes or more complex inputs. This information would be valuable for configuring computation resources effectively.

Thanks again for your helpful advice!

Yes, the easiest method is to decrease the batch_size. Another simple way is cutting protein pockets, such as retaining only residues within a distance of 10A from the center of the pocket of interest.