take care of data collation, fix some bugs due to inplace ops

Author: lucidrains

alphafold3_pytorch/__init__.py

@@ -34,6 +34,7 @@
 
 from alphafold3_pytorch.trainer import (
     Trainer,
+    DataLoader,
     Alphafold3Input
 )
 

alphafold3_pytorch/alphafold3.py

@@ -212,7 +212,7 @@ def repeat_consecutive_with_lens(
 
     output_indices = torch.zeros((batch, max_len + 1), device = device, dtype = torch.long)
 
-    indices.masked_fill_(~mask, max_len) # scatter to sink position for padding
+    indices = indices.masked_fill(~mask, max_len) # scatter to sink position for padding
     indices = rearrange(indices, 'b n w -> b (n w)')
 
     # scatter
@@ -3062,6 +3062,15 @@ def forward(
 
         atom_seq_len = atom_inputs.shape[-2]
 
+        # soft validate
+
+        valid_atom_len_mask = residue_atom_lens >= 0
+
+        residue_atom_lens = residue_atom_lens.masked_fill(~valid_atom_len_mask, 0)
+        residue_atom_indices = residue_atom_indices.masked_fill(~valid_atom_len_mask, 0)
+
+        assert (residue_atom_indices < residue_atom_lens)[valid_atom_len_mask].all(), 'residue_atom_indices cannot have an index that exceeds the length of the atoms for that residue as given by residue_atom_lens'
+
         assert exists(residue_atom_lens) or exists(atom_mask)
 
         # if atompair inputs are not windowed, window it
@@ -3079,7 +3088,7 @@ def forward(
         # handle offsets for residue atom indices
 
         if exists(residue_atom_indices):
-            residue_atom_indices += F.pad(residue_atom_lens, (-1, 1), value = 0)
+            residue_atom_indices = residue_atom_indices + F.pad(residue_atom_lens, (-1, 1), value = 0)
 
         # get atom sequence length and residue sequence length depending on whether using packed atomic seq
 
@@ -3118,7 +3127,7 @@ def forward(
 
             token_bond = token_bond | rearrange(token_bond, 'b i j -> b j i')
             diagonal = torch.eye(seq_len, device = self.device, dtype = torch.bool)
-            token_bond.masked_fill_(diagonal, False)
+            token_bond = token_bond.masked_fill(diagonal, False)
         else:
             seq_arange = torch.arange(seq_len, device = self.device)
             token_bond = einx.subtract('i, j -> i j', seq_arange, seq_arange).abs() == 1

alphafold3_pytorch/trainer.py

@@ -3,16 +3,19 @@
 from pathlib import Path
 
 from alphafold3_pytorch.alphafold3 import Alphafold3
+from alphafold3_pytorch.attention import pad_at_dim
 
-from typing import TypedDict
+from typing import TypedDict, List
 from alphafold3_pytorch.typing import (
     typecheck,
     Int, Bool, Float
 )
 
 import torch
+from torch import Tensor
 from torch.optim import Adam, Optimizer
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import Dataset, DataLoader as OrigDataLoader
+from torch.nn.utils.rnn import pad_sequence
 from torch.optim.lr_scheduler import LambdaLR, LRScheduler
 
 from ema_pytorch import EMA
@@ -24,7 +27,7 @@
 @typecheck
 class Alphafold3Input(TypedDict):
     atom_inputs:                Float['m dai']
-    residue_atom_lens:          Int['n 2']
+    residue_atom_lens:          Int[' n']
     atompair_inputs:            Float['m m dapi'] | Float['nw w (w*2) dapi']
     additional_residue_feats:   Float['n 10']
     templates:                  Float['t n n dt']
@@ -70,6 +73,62 @@ def accum_dict(
 
     return past_losses
 
+# dataloader and collation fn
+
+@typecheck
+def collate_af3_inputs(
+    inputs: List[Alphafold3Input],
+    int_pad_value = -1
+):
+    # separate input dictionary into keys and values
+
+    keys = inputs[0].keys()
+    inputs = [i.values() for i in inputs]
+
+    outputs = []
+
+    for grouped in zip(*inputs):
+        # if all None, just return None
+
+        if not any([*map(exists, grouped)]):
+            outputs.append(None)
+            continue
+
+        # use -1 for padding int values, for assuming int are labels - if not, handle within alphafold3
+
+        pad_value = int_pad_value if grouped[0].dtype in (torch.int, torch.long) else 0
+
+        # get the max lengths across all dimensions
+
+        shapes_as_tensor = torch.stack([Tensor(tuple(g.shape)) for g in grouped], dim = -1)
+
+        max_lengths = shapes_as_tensor.int().amax(dim = -1)
+
+        # pad across all dimensions
+
+        padded_inputs = []
+
+        for inp in grouped:
+            for dim, max_length in enumerate(max_lengths.tolist()):
+                inp = pad_at_dim(inp, (0, max_length - inp.shape[dim]), value = pad_value, dim = dim)
+
+            padded_inputs.append(inp)
+
+        # stack
+
+        stacked = torch.stack(padded_inputs)
+
+        outputs.append(stacked)
+
+    # reconstitute dictionary
+
+    return dict(tuple(zip(keys, outputs)))
+
+def DataLoader(*args, **kwargs):
+    return OrigDataLoader(*args, collate_fn = collate_af3_inputs, **kwargs)
+
+# default scheduler used in paper w/ warmup
+
 def default_lambda_lr_fn(steps):
     # 1000 step warmup
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.1.5"
+version = "0.1.6"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

tests/test_trainer.py

@@ -2,14 +2,16 @@
 os.environ['TYPECHECK'] = 'True'
 
 from pathlib import Path
+from random import randrange
 
 import pytest
 import torch
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import Dataset
 
 from alphafold3_pytorch import (
     Alphafold3,
     Alphafold3Input,
+    DataLoader,
     Trainer
 )
 
@@ -19,25 +21,24 @@ class MockAtomDataset(Dataset):
     def __init__(
         self,
         data_length,
-        seq_len = 16,
+        max_seq_len = 16,
         atoms_per_window = 4
     ):
         self.data_length = data_length
-        self.seq_len = seq_len
+        self.max_seq_len = max_seq_len
         self.atoms_per_window = atoms_per_window
-        self.atom_seq_len = seq_len * atoms_per_window
 
     def __len__(self):
         return self.data_length
 
     def __getitem__(self, idx):
-        seq_len = self.seq_len
-        atom_seq_len = self.atom_seq_len
+        seq_len = randrange(1, self.max_seq_len)
+        atom_seq_len = self.atoms_per_window * seq_len
 
         atom_inputs = torch.randn(atom_seq_len, 77)
         atompair_inputs = torch.randn(atom_seq_len, atom_seq_len, 5)
 
-        residue_atom_lens = torch.randint(0, self.atoms_per_window, (seq_len,))
+        residue_atom_lens = torch.randint(1, self.atoms_per_window, (seq_len,))
         additional_residue_feats = torch.randn(seq_len, 10)
 
         templates = torch.randn(2, seq_len, seq_len, 44)