LocScale-SURFER bundle development for Mac OS

[A.Bharadwaj@…]

Hello,

I'm developing the bundle LocScale-SURFER with Arjen Jakobi.
  https://cxtoolshed.rbvi.ucsf.edu/apps/chimeraxlocscalesurfer

The tool uses a pytorch model to predict detergent micelles for the cryoEM maps of membrane proteins. To speed up the computation, we use a GPU for the predictions. Currently it does not use GPUs for Apple MPS systems. One problem I face while implementing mps support is that  Conv3D layers are not supported on MPS. Since our model uses Conv3D layers, this really prevents us from using the GPUs in Apple devices.

However, I notice that if I run the prediction outside of ChimeraX development environment and inside a custom conda environment I am able to run the pytorch model with MPS support. I installed pytorch 2.7 using defaults conda channels. I wondered if it is possible to install packages from conda channels for Bundle development and distribution?

Note;
Currently I am testing on Apple silicon M2, ChimeraX version 1.8

Let me know if I need to provide anything else.

Cheers,
Alok

[pett]

Reported by Alok Bharadwaj

[Tom Goddard]

Hi Alok,

I have never seen Conda packages installed into ChimeraX. But the PyPi torch is version 2.7.0 so I think it is no different than the Conda torch. The difference is probably that you can tell torch to fallback to CPU computation for operations that are not available in MPS. You do this by setting environment variable PYTORCH_ENABLE_MPS_FALLBACK = 1. I use this in the ChimeraX Boltz structure prediction tool which also uses torch. If your Conv3D layers are taking most of the time then probably MPS won't speed things up much if it has to fallback to CPU. I think you can set this environment variable in Python.

import os
os.environPYTORCH_ENABLE_MPS_FALLBACK = '1'

I am not sure if that will work. In the ChimeraX Boltz tool I use pytorch in a subprocess rather than running directly in ChimeraX.

Tom

[A.Bharadwaj@…]

Thanks for the response, Tom!

Setting the PYTORCH_ENABLE_MPS_FALLBACK = 1 helped overcome the problem I had with Conv3D using torch. However that did not work for a different layer of the neural network (conv3dtranspose). This is apparently a bug in PyTorch code and I have raised an issue with the dev team regarding that. It will be fixed in the next nightly release of pytorch.
​https://github.com/pytorch/pytorch/issues/154615

For the moment, I will not make use of MPS for running the prediction. So locscalesurfer will run slower on Apple machines until the bug is fixed.

Thanks for the help!

Cheers,
Alok

[Tom Goddard]

Thanks Alok! It is good to contribute the bug report to PyTorch since that project is so important to machine learning in biology. Even when they fix the CPU fallback of conv3dtranspose, it may be that much of your network is using the CPU instead of Metal GPU shaders so it might not see much improvement. But I hope PyTorch MPS will improve so everything can run on the Mac GPU. There is a PyTorch issue where people note the missing torch functions they want MPS to support

​https://github.com/pytorch/pytorch/issues/77764

Perhaps Conv3d is already there.