Some questions regarding reducing the size of both input and model? about gemnet_pytorch HOT 3 CLOSED

Thank you for the interest in our work.
Reducing the size of the model input can only be achieved by decreasing the two cutoff distances. As discussed in the appendix of our paper this reduces the prediction accuracy. However, this is a tradeoff with computational resources so you could give this a shot.
Further, to reduce the model size this can be achieved by reducing the embedding sizes (emb_size_...) or the number of interaction blocks (num_blocks). Another approach can be to use GemNet-T instead of GemNet-Q (as done for OC20) or even use the direct force prediction models GemNet-dQ/T. This all significantly lowers the memory consumption but again is a tradeoff with prediction accuracy.
Let me know if this helps.

from gemnet_pytorch.

Moreover, I also need to reduce the model size. After I turned the batch size of 32 into 2, it becomes another problem of utilizing GemNet to process that. Can you give me some hints on getting a smaller model?

from gemnet_pytorch.

I second the points Flo mentioned. I think your very first step should be using GemNet-T instead of GemNet-Q. That will already give a huge improvement. I wouldn't recommend reducing cutoffs below 4A.

I would not recommend using direct models (GemNet-dT) for molecular dynamics, since they can lead to unstable trajectories. If your task is not simulation, though, then you can use a direct model to gain another 2-3x reduction in memory and runtime.

Embedding sizes and depth (num_blocks) would then be the third thing to consider. It's hard to give advice on which value presents the best trade-off, since they will likely all reduce accuracy. A first step might be to simply half all embedding sizes and use 3 blocks instead of 4. But probably you should not reduce anything below 16. Note that emb_size_quad, emb_size_sbf and emb_size_bil_quad are not relevant for GemNet-T. You can then increase embeddings again to see which ones are the most important for your application.

Overall, I think it should easily be possible to find well-performing settings for GemNet-T on proteins that fit in 45GiB.

from gemnet_pytorch.