v2.3.4 - fix memory leaks (attempt 2) (google-deepmind#6)

* fix memory leaks

various edits to fix memory leaks
memory leak fix

* v2.3.4 - fix memory leaks

another attempt to fix memory leaks!

* Update config.py

* bugfix - num-ensemble

alphafold/common/confidence.py

@@ -17,8 +17,20 @@
 import jax.numpy as jnp
 import jax
 import numpy as np
+from alphafold.common import residue_constants
 import scipy.special
 
+def compute_tol(prev_pos, current_pos, mask, use_jnp=False):
+ # Early stopping criteria based on criteria used in
+ # AF2Complex: https://www.nature.com/articles/s41467-022-29394-2 
+ _np = jnp if use_jnp else np
+ dist = lambda x:_np.sqrt(((x[:,None] - x[None,:])**2).sum(-1))
+ ca_idx = residue_constants.atom_order['CA']
+ sq_diff = _np.square(dist(prev_pos[:,ca_idx])-dist(current_pos[:,ca_idx]))
+ mask_2d = mask[:,None] * mask[None,:]
+ return _np.sqrt((sq_diff * mask_2d).sum()/mask_2d.sum() + 1e-8)
+
+
 def compute_plddt(logits, use_jnp=False):
  """Computes per-residue pLDDT from logits.
  Args:

alphafold/model/config.py

@@ -22,7 +22,6 @@
 NUM_EXTRA_SEQ = shape_placeholders.NUM_EXTRA_SEQ
 NUM_TEMPLATES = shape_placeholders.NUM_TEMPLATES
 
-
 def model_config(name: str) -> ml_collections.ConfigDict:
  """Get the ConfigDict of a CASP14 model."""
 
@@ -378,7 +377,7 @@ def model_config(name: str) -> ml_collections.ConfigDict:
  },
  'global_config': {
  'bfloat16': True,
- 'bfloat16_output': True,
+ 'bfloat16_output': False,
  'deterministic': False,
  'multimer_mode': False,
  'subbatch_size': 4,
@@ -616,7 +615,7 @@ def model_config(name: str) -> ml_collections.ConfigDict:
  },
  'global_config': {
  'bfloat16': True,
- 'bfloat16_output': True,
+ 'bfloat16_output': False,
  'deterministic': False,
  'multimer_mode': True,
  'subbatch_size': 4,

alphafold/model/model.py

@@ -148,25 +148,30 @@ def predict(self,
  L = aatype.shape[1]
 
  # initialize
- def z(shape, dtype=np.float32): return np.zeros(shape, dtype=dtype)
- dtype = jnp.bfloat16 if self.config.model.global_config.bfloat16 else np.float32
- prev = {'prev_msa_first_row': z([L,256], dtype),
- 'prev_pair': z([L,L,128], dtype),
- 'prev_pos': z([L,37,3])}
+
+ zeros = lambda shape: np.zeros(shape, dtype=np.float16)
+ prev = {'prev_msa_first_row': zeros([L,256]),
+ 'prev_pair': zeros([L,L,128]),
+ 'prev_pos': zeros([L,37,3])}
 
  def run(key, feat, prev):
- outputs = jax.tree_map(lambda x:np.asarray(x),
- self.apply(self.params, key, {**feat, "prev":prev}))
- prev = outputs.pop("prev")
- return outputs, prev
+ def _jnp_to_np(x):
+ for k, v in x.items():
+ if isinstance(v, dict):
+ x[k] = _jnp_to_np(v)
+ else:
+ x[k] = np.asarray(v,np.float16)
+ return x
+ result = _jnp_to_np(self.apply(self.params, key, {**feat, "prev":prev}))
+ prev = result.pop("prev")
+ return result, prev
+
 
  # initialize random key
  key = jax.random.PRNGKey(random_seed)
 
  # iterate through recyckes
- stop = False
- for r in range(num_iters):
-
+ for r in range(num_iters): 
  # grab subset of features
  if self.multimer_mode:
  sub_feat = feat
@@ -180,30 +185,18 @@ def run(key, feat, prev):
  result, prev = run(sub_key, sub_feat, prev)
 
  if return_representations:
- result["representations"] = {"pair": prev["prev_pair"].astype(np.float32),
- "single": prev["prev_msa_first_row"].astype(np.float32)}
- # decide when to stop
- tol = self.config.model.recycle_early_stop_tolerance
- sco = self.config.model.stop_at_score
- if result["ranking_confidence"] > sco:
- stop = True
- if not stop and tol > 0:
- ca_idx = residue_constants.atom_order['CA']
- pos = result["structure_module"]["final_atom_positions"][:,ca_idx]
- dist = np.sqrt(np.square(pos[:,None]-pos[None,:]).sum(-1))
- if r > 0:
- sq_diff = np.square(dist - prev_dist)
- seq_mask = feat["seq_mask"] if self.multimer_mode else feat["seq_mask"][0]
- mask_2d = seq_mask[:,None] * seq_mask[None,:]
- result["diff"] = np.sqrt((sq_diff * mask_2d).sum()/mask_2d.sum())
- if result["diff"] < tol:
- stop = True
- prev_dist = dist
 
+ result["representations"] = {"pair": prev["prev_pair"],
+ "single": prev["prev_msa_first_row"]}
+
  # callback
  if callback is not None: callback(result, r)
 
- if stop: break
+ # decide when to stop
+ if result["ranking_confidence"] > self.config.model.stop_at_score:
+ break
+ if r > 0 and result["tol"] < self.config.model.recycle_early_stop_tolerance:
+ break
 
  logging.info('Output shape was %s', tree.map_structure(lambda x: x.shape, result))
  return result, r

alphafold/model/modules.py

@@ -182,12 +182,19 @@ def get_prev(ret):
  prev = batch.pop("prev",None)
  if batch["aatype"].ndim == 2:
  batch = jax.tree_map(lambda x:x[0], batch)
+
+ # initialize
  if prev is None:
- L = batch["aatype"].shape[0]
- dtype = jnp.bfloat16 if self.global_config.bfloat16 else jnp.float32
- prev = {'prev_msa_first_row': jnp.zeros([L,256], dtype=dtype),
- 'prev_pair': jnp.zeros([L,L,128],dtype=dtype),
+
+ L = num_residues
+ prev = {'prev_msa_first_row': jnp.zeros([L,256]),
+ 'prev_pair': jnp.zeros([L,L,128]),
  'prev_pos': jnp.zeros([L,37,3])}
+ else:
+ for k,v in prev.items():
+ if v.dtype == jnp.float16:
+ prev[k] = v.astype(jnp.float32)
+
 
  ret = impl(batch={**batch, **prev}, is_training=is_training)
  ret["prev"] = get_prev(ret)
@@ -200,7 +207,13 @@ def get_prev(ret):
  mask=batch["seq_mask"],
  rank_by=self.config.rank_by,
  use_jnp=True))
-
+
+ ret["tol"] = confidence.compute_tol(
+ prev["prev_pos"], 
+ ret["prev"]["prev_pos"],
+ batch["seq_mask"], 
+ use_jnp=True)
+
  return ret
 
 class AlphaFoldIteration(hk.Module):
@@ -426,13 +439,18 @@ def slice_recycle_idx(x):
  ensemble_representations=ensemble_representations)
 
  emb_config = self.config.embeddings_and_evoformer
- prev = batch.pop("prev", None)
+ # initialize
+ prev = batch.pop("prev", None) 
  if prev is None:
  L = num_residues
- dtype = jnp.bfloat16 if self.global_config.bfloat16 else jnp.float32
- prev = {'prev_msa_first_row': jnp.zeros([L,256], dtype=dtype),
- 'prev_pair': jnp.zeros([L,L,128],dtype=dtype),
+ prev = {'prev_msa_first_row': jnp.zeros([L,256]),
+ 'prev_pair': jnp.zeros([L,L,128]),
  'prev_pos': jnp.zeros([L,37,3])}
+ else:
+ for k,v in prev.items():
+ if v.dtype == jnp.float16:
+ prev[k] = v.astype(jnp.float32)
+
 
  ret = do_call(prev=prev, recycle_idx=0)
  ret["prev"] = get_prev(ret)
@@ -446,9 +464,15 @@ def slice_recycle_idx(x):
  # add confidence metrics
  ret.update(confidence.get_confidence_metrics(
  prediction_result=ret,
- mask=batch["seq_mask"],
+ mask=batch["seq_mask"][0],
  rank_by=self.config.rank_by,
- use_jnp=True)) 
+ use_jnp=True))
+
+ ret["tol"] = confidence.compute_tol(
+ prev["prev_pos"], 
+ ret["prev"]["prev_pos"],
+ batch["seq_mask"][0], 
+ use_jnp=True)
 
  return ret
 

alphafold/model/modules_multimer.py

@@ -442,8 +442,20 @@ def apply_network(prev, safe_key):
  batch=recycled_batch,
  is_training=is_training,
  safe_key=safe_key)
-
- ret = apply_network(prev=batch.pop("prev"), safe_key=safe_key)
+
+ # initialize
+ prev = batch.pop("prev", None) 
+ if prev is None:
+ L = num_residues
+ prev = {'prev_msa_first_row': jnp.zeros([L,256]),
+ 'prev_pair': jnp.zeros([L,L,128]),
+ 'prev_pos': jnp.zeros([L,37,3])}
+ else:
+ for k,v in prev.items():
+ if v.dtype == jnp.float16:
+ prev[k] = v.astype(jnp.float32)
+
+ ret = apply_network(prev=prev, safe_key=safe_key)
  ret["prev"] = get_prev(ret)
 
  if not return_representations:
@@ -456,6 +468,12 @@ def apply_network(prev, safe_key):
  rank_by=self.config.rank_by,
  use_jnp=True))
 
+ ret["tol"] = confidence.compute_tol(
+ prev["prev_pos"], 
+ ret["prev"]["prev_pos"],
+ batch["seq_mask"], 
+ use_jnp=True)
+
  return ret
 
 class EmbeddingsAndEvoformer(hk.Module):

setup.py

@@ -18,7 +18,7 @@
 
 setup(
  name='alphafold-colabfold',
- version='2.3.3',
+ version='2.3.4',
  long_description_content_type='text/markdown',
  description='An implementation of the inference pipeline of AlphaFold v2.3.1. '
  'This is a completely new model that was entered as AlphaFold2 in CASP14 '