Merge pull request #375 from bigict/xpu

refactor: code reformat

Author: chungongyu

profold2/command/evaluator.py

@@ -70,7 +70,7 @@ def data_eval(idx, batch):
     # predict - out is (batch, L * 3, 3)
     with timing(f'Running model on {fasta_name} {fasta_len}', logging.debug):
       with torch.no_grad():
-        with worker.autocast_ctx(args.amp_enabled):
+        with accelerator.amp(args.amp_enabled):  # Automatic Mixed Precision
           r = ReturnValues(
               **model(
                   batch=batch,  # pylint: disable=not-callable

profold2/command/predictor.py

@@ -19,7 +19,7 @@
 from profold2.data import dataset
 from profold2.data.dataset import ProteinSequenceDataset
 from profold2.data.parsers import parse_fasta
-from profold2.data.utils import parse_seq_index, pdb_from_prediction, str_seq_index
+from profold2.data.utils import pdb_from_prediction
 from profold2.model import accelerator, profiler, snapshot, FeatureBuilder, ReturnValues
 from profold2.utils import exists, timing
 
@@ -222,7 +222,7 @@ def process_structure(model_name, pdb_str, plddt=None):
                   timing_callback, timings, f'predict_{model_name}'
               )
           ):
-            with worker.autocast_ctx(args.amp_enabled):
+            with accelerator.amp(args.amp_enabled):  # Automatic Mixed Precision
               r = ReturnValues(
                   **model(
                       batch=feats,

profold2/command/trainer.py

@@ -352,7 +352,7 @@ def _step(data_loader, it, writer, stage='train', batch_callback=None):
       with no_sync_ctx(
           it != global_step and jt + 1 != args.gradient_accumulate_every, model
       ):
-        with worker.autocast_ctx(grad_scaler.is_enabled()):
+        with accelerator.amp(grad_scaler.is_enabled()):  # Automatic Mixed Precision
           r = ReturnValues(
               **model(
                   batch=batch,

profold2/command/worker.py

@@ -1,7 +1,6 @@
 """Wrap distibuted env
 """
 import os
-import contextlib
 import functools
 import logging
 from logging.handlers import QueueHandler, QueueListener
@@ -14,24 +13,7 @@
 
 from profold2.model import accelerator, AlphaFold2
 from profold2.model.commons import torch_allow_tf32
-from profold2.utils import default, env, exists, version_cmp
-
-
-@contextlib.contextmanager
-def autocast_ctx(cond):
-  if cond:
-    dtype = torch.float16
-    if hasattr(torch.cuda, 'is_bf16_supported'):
-      if torch.cuda.is_bf16_supported():
-        dtype = torch.bfloat16
-    ctx = functools.partial(accelerator.autocast, dtype=dtype)
-    # FIXED ME: cache_enabled=True will crash :(
-    if version_cmp(torch.__version__, '1.10.0') >= 0:
-      ctx = functools.partial(ctx, cache_enabled=False)
-    with ctx():
-      yield
-  else:
-    yield
+from profold2.utils import exists, version_cmp
 
 
 class _WorkerLogRecordFactory(object):

profold2/model/accelerator.py

@@ -1,29 +1,51 @@
 """Wrap GPU, MLU etc
 """
+import contextlib
 from datetime import timedelta
 import functools
 import logging
 from typing import Optional
 
 import torch
 
-from profold2.utils import default, env
+from profold2.utils import default, env, version_cmp
 
 
 def device_count():
   return torch.cuda.device_count()
 
+
 def device_type():
   if torch.cuda.is_available():
     return 'cuda'
   return 'cpu'
 
+
 def world_size(nnodes: Optional[int] = None):
   return env('WORLD_SIZE', defval=device_count() * default(nnodes, 1), dtype=int)
 
+
 autocast = functools.partial(torch.amp.autocast, device_type())
 GradScaler = functools.partial(torch.amp.GradScaler, device_type())
 
+
+@contextlib.contextmanager
+def amp(enabled: bool = True):
+  if enabled:
+    dtype = torch.float16
+    if hasattr(torch.cuda, 'is_bf16_supported'):
+      if torch.cuda.is_bf16_supported():
+        dtype = torch.bfloat16
+    ctx = functools.partial(autocast, dtype=dtype)
+    # FIXED ME: cache_enabled=True will crash :(
+    if version_cmp(torch.__version__, '1.10.0') >= 0:
+      ctx = functools.partial(ctx, cache_enabled=False)
+    with ctx():
+      yield
+  else:
+    yield
+
+
 class XPU(object):
   """Wrap distibuted XPU(GPU,MLU etc)
   """
@@ -51,9 +73,7 @@ def device(self):
   @property
   def rank(self):
     return env(
-        'RANK',
-        defval=device_count() * self.node_rank + self.local_rank,
-        dtype=int
+        'RANK', defval=device_count() * self.node_rank + self.local_rank, dtype=int
     )
 
   def memory_summary(self):