Merge pull request #19 from RuiWang1998/main

fixed a bug where user would have to fix the subbatch size

Forcing this through as it is a bit urgent and it takes a while for others to respond. It seems to run just fine on my end.

omegafold/modules.py

@@ -522,6 +522,24 @@ def forward(
         return attn_out
 
 
+def _get_sharded_stacked(
+        edge_repr: torch.Tensor,
+        subbatch_size: int
+):
+    subbatch_size = subbatch_size or edge_repr.shape[-2]
+    idx = 0
+    start, end = 0, subbatch_size
+    while start < edge_repr.shape[-2]:
+        yield start, end, torch.stack(
+            [
+                edge_repr[start:end],
+                edge_repr.transpose(-2, -3)[start:end]
+            ], dim=-1
+        )
+        idx += 1
+        start, end = idx * subbatch_size, (idx + 1) * subbatch_size
+
+
 class GeometricAttention(OFModule):
     """We have a lot of stuff here for GRAM reduction
 
@@ -577,37 +595,20 @@ def _get_attended(
             device=edge_repr.device
         )
         b += utils.mask2bias(mask)
-        for s, e, edge_r in self._get_sharded_stacked(
+        for s, e, edge_r in _get_sharded_stacked(
                 edge_repr, subbatch_size=fwd_cfg.subbatch_size
         ):
             b[..., s:e, :] = torch.einsum(
                 '...qkcr,crh->...rhqk', edge_r, self.linear_b_weights
             ) + self.linear_b_bias
-        for s, e, edge_r in self._get_sharded_stacked(
+        for s, e, edge_r in _get_sharded_stacked(
                 edge_repr, subbatch_size=fwd_cfg.subbatch_size
         ):
             attended[s:e] = self.attention(
                 edge_r, edge_r, b, fwd_cfg=fwd_cfg
             )
         return attended[..., 0] + attended[..., 1].transpose(-2, -3)
 
-    def _get_sharded_stacked(
-            self,
-            edge_repr: torch.Tensor,
-            subbatch_size: int
-    ):
-        idx = 0
-        start, end = 0, subbatch_size
-        while start < edge_repr.shape[-2]:
-            yield start, end, torch.stack(
-                [
-                    edge_repr[start:end],
-                    edge_repr.transpose(-2, -3)[start:end]
-                ], dim=-1
-            )
-            idx += 1
-            start, end = idx * subbatch_size, (idx + 1) * subbatch_size
-
     def _get_gated(self, edge_repr: torch.Tensor, mask: torch.Tensor, fwd_cfg):
         gated = torch.empty(
             *edge_repr.shape[:2],
@@ -616,7 +617,7 @@ def _get_gated(self, edge_repr: torch.Tensor, mask: torch.Tensor, fwd_cfg):
             device=edge_repr.device,
             dtype=edge_repr.dtype
         )
-        for s_row, e_row, edge_row in self._get_sharded_stacked(
+        for s_row, e_row, edge_row in _get_sharded_stacked(
                 edge_repr, subbatch_size=fwd_cfg.subbatch_size
         ):
             act_row = self._get_act_row(edge_row, mask[s_row:e_row])
@@ -627,7 +628,7 @@ def _get_gated(self, edge_repr: torch.Tensor, mask: torch.Tensor, fwd_cfg):
                     self.act_w[..., -self.d_edge:]
                 ) + self.act_b[..., -self.d_edge:]
             )
-            for s_col, e_col, edge_col, in self._get_sharded_stacked(
+            for s_col, e_col, edge_col, in _get_sharded_stacked(
                     edge_repr, subbatch_size=fwd_cfg.subbatch_size
             ):
                 act_col = self._get_act_col(edge_col, mask[s_col:e_col])

omegafold/pipeline.py

@@ -14,28 +14,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # =============================================================================
-from __future__ import annotations
 """
 This file contains the utilities that we use for the entire inference pipeline
 """
 # =============================================================================
 # Imports
 # =============================================================================
+from __future__ import annotations
+
 import argparse
 import collections
 import logging
 import ntpath
 import os
 import os.path
 import pathlib
-import sys
 import typing
 
 from Bio import PDB as PDB
 from Bio.PDB import StructureBuilder
 import torch
 from torch import hub
-from torch.backends import cuda, cudnn
+from torch.backends import cuda, cudnn, mps
 from torch.utils.hipify import hipify_python
 
 from omegafold import utils
@@ -259,6 +259,39 @@ def _load_weights(
     return torch.load(weights_file, map_location='cpu')
 
 
+def _get_device(device) -> str:
+    """
+    Infer the accelerator
+
+    Args:
+        device: the device type
+
+    Returns:
+
+    """
+    if device is None:
+        if torch.cuda.is_available():
+            return "cuda"
+        elif mps.is_available():
+            return "mps"
+        else:
+            return 'cpu'
+    elif device == 'cpu':
+        return device
+    elif device.startswith('cuda'):
+        if torch.cuda.is_available():
+            return device
+        else:
+            raise ValueError(f"Device cuda is not available")
+    elif device == "mps":
+        if mps.is_available():
+            return device
+        else:
+            raise ValueError(f"Device mps is not available")
+    else:
+        raise ValueError(f"Device type {device} is not available")
+
+
 def get_args() -> typing.Tuple[
     argparse.Namespace, collections.OrderedDict, argparse.Namespace]:
     """
@@ -314,8 +347,10 @@ def get_args() -> typing.Tuple[
         """
     )
     parser.add_argument(
-        '--device', default='cuda', type=str,
-        help='The device on which the model will be running, default to cuda'
+        '--device', default=None, type=str,
+        help=
+        'The device on which the model will be running, '
+        'default to the accelerator that we can find'
     )
     parser.add_argument(
         '--weights_file',
@@ -360,6 +395,8 @@ def get_args() -> typing.Tuple[
         num_recycle=args.num_cycle,
     )
 
+    args.device = _get_device(args.device)
+
     return args, weights, forward_config