From ae9ea4784b47369a7337c669a9e39a168b63cfea Mon Sep 17 00:00:00 2001 From: fnachon Date: Tue, 1 Apr 2025 23:15:33 +0200 Subject: [PATCH 1/4] Add mps device for Apple silicon gpu --- README.md | 10 ++++++++++ environment_Mac.yaml | 30 ++++++++++++++++++++++++++++++ src/data/so3_utils.py | 2 +- src/generate.py | 13 +++++++++++-- src/model/gvp_transformer.py | 2 +- 5 files changed, 53 insertions(+), 4 deletions(-) create mode 100644 environment_Mac.yaml diff --git a/README.md b/README.md index fbaffed..dfefc78 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,16 @@ and add the Gnina executable for docking score computation wget https://github.com/gnina/gnina/releases/download/v1.1/gnina -O $CONDA_PREFIX/bin/gnina chmod +x $CONDA_PREFIX/bin/gnina ``` +### Conda Environment for Mac (Apple Silicon) + +Create a conda/mamba environment +``` +conda env create -f environment_Mac.yaml -n drugflow +conda activate drugflow +``` + +Gnina is not available for Apple Silicon as it depends on libmolgrid, which heavily depends on cuda kernels. Qvina is installed in the conda environment as an alternative for docking scoring. + ### Docker Container diff --git a/environment_Mac.yaml b/environment_Mac.yaml new file mode 100644 index 0000000..b359742 --- /dev/null +++ b/environment_Mac.yaml @@ -0,0 +1,30 @@ +name: drugflow2 + +channels: + - conda-forge + - pyg + - bioconda + +dependencies: + - python=3.11.11 + - pytorch=2.4.1 + - pytorch-lightning=2.5.0 + - rdkit=2024.03.5 + - openbabel=3.1.1 + - biopython=1.85 + - scipy=1.15.2 + - pyg=2.6.1 + - pytorch-scatter=2.1.2 + - ProDy=2.4.1 + - wandb=0.16.6 + - pandas=2.2.2 + - posebusters=0.3.1 + - fcd=1.2.2 + - webdataset=0.2.100 + - prolif=2.0.3 + - reduce=4.14 + - qvina=2.1.0 + - pip=25.0.1 + - pip: + - useful_rdkit_utils==0.65 + diff --git a/src/data/so3_utils.py b/src/data/so3_utils.py index 0d10c99..f60a24b 100644 --- a/src/data/so3_utils.py +++ b/src/data/so3_utils.py @@ -344,7 +344,7 @@ def log_not_from_identity(point, base_point): # decorator if torch.__version__ >= '2.0.0': - GEOMSTATS_DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' + GEOMSTATS_DEVICE = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu' def geomstats_tensor_type(func): def inner(*args, **kwargs): diff --git a/src/generate.py b/src/generate.py index 9bfa7c3..5e2bf4f 100644 --- a/src/generate.py +++ b/src/generate.py @@ -4,6 +4,7 @@ import warnings import tempfile import pandas as pd +import torch from Bio.PDB import PDBParser from pathlib import Path @@ -66,7 +67,7 @@ def aggregate_metrics(table): p.add_argument('--batch_size', type=int, required=False, default=32, help="Batch size.") p.add_argument('--pocket_distance_cutoff', type=float, required=False, default=8.0, help="Distance cutoff to define the pocket around the reference ligand.") p.add_argument('--n_steps', type=int, required=False, default=None, help="Number of denoising steps.") - p.add_argument('--device', type=str, required=False, default='cuda:0', help="Device to use.") + p.add_argument('--device', type=str, required=False, default='cuda:0', help="Device to use (mps, cuda:0).") p.add_argument('--datadir', type=Path, required=False, default=Path(basedir, 'src', 'default'), help="Needs to be specified to sample molecule sizes.") p.add_argument('--seed', type=int, required=False, default=42, help="Random seed.") p.add_argument('--filter', action='store_true', required=False, default=False, help="Apply basic filters and keep sampling until `n_samples` molecules passing these filters are found.") @@ -86,11 +87,19 @@ def aggregate_metrics(table): if not args.filter: args.batch_size = min(args.batch_size, args.n_samples) + + # Select mps if available (Apple Silicon) + if torch.backends.mps.is_available(): + args.device = 'mps' # Loading model chkpt_path = Path(args.checkpoint) chkpt_name = chkpt_path.parts[-1].split('.')[0] - model = DrugFlow.load_from_checkpoint(args.checkpoint, map_location=args.device, strict=False) + if torch.backends.mps.is_available(): + model = DrugFlow.load_from_checkpoint(args.checkpoint, map_location='cpu', strict=False) # load model to cpu + model = model.to(args.device, dtype=torch.float32) # move model to mps with float32 type + else: + model = DrugFlow.load_from_checkpoint(args.checkpoint, map_location=args.device, strict=False) if args.datadir is not None: model.datadir = args.datadir diff --git a/src/model/gvp_transformer.py b/src/model/gvp_transformer.py index f36bdc9..6fa226b 100644 --- a/src/model/gvp_transformer.py +++ b/src/model/gvp_transformer.py @@ -454,7 +454,7 @@ def test_equivariance(model, nodes, edges, glob_feat): db = 10 attn_heads = 9 - device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + device = torch.device('mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu') nodes = randn(n_nodes, node_dim, device=device) From 11e052ca777a57b393b50b1de96f2d8efb501b9b Mon Sep 17 00:00:00 2001 From: fnachon Date: Tue, 1 Apr 2025 23:49:04 +0200 Subject: [PATCH 2/4] Update .gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index c0d4110..5529aa5 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ notebooks/ *DS_Store samples/ crossdocked/ +/checkpoints From 96ae900d9e4c223621b712fc6123a10dc8e7305c Mon Sep 17 00:00:00 2001 From: fnachon Date: Tue, 1 Apr 2025 23:54:14 +0200 Subject: [PATCH 3/4] Update .gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 5529aa5..a2d5269 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,4 @@ notebooks/ samples/ crossdocked/ /checkpoints +DrugFlow on Apple Silicon.md From 395b88679995766db7fbd17214ed34ff49c3e13f Mon Sep 17 00:00:00 2001 From: fnachon Date: Tue, 1 Apr 2025 23:57:37 +0200 Subject: [PATCH 4/4] Update .gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index a2d5269..60c4d97 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ samples/ crossdocked/ /checkpoints DrugFlow on Apple Silicon.md +/output