alphagenome-pytorch 0.0.8__tar.gz → 0.0.10__tar.gz

{alphagenome_pytorch-0.0.8 → alphagenome_pytorch-0.0.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alphagenome-pytorch
-Version: 0.0.8
+Version: 0.0.10
 Summary: AlphaGenome
 Project-URL: Homepage, https://pypi.org/project/alphagenome-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/alphagenome
@@ -58,13 +58,13 @@ $ pip install alphagenome-pytorch
 
 ```python
 import torch
-from alphagenome_pytorch import TransformerTower
+from alphagenome_pytorch import AlphaGenome
 
-transformer = TransformerTower(dim = 768, dim_pairwise = 128)
+model = AlphaGenome()
 
-single = torch.randn(2, 512, 768)
+dna = torch.randint(0, 5, (2, 8192))
 
-attended_single, attended_pairwise = transformer(single)
+pred_nucleotide, single, pairwise = model(dna) # (2, 8192, 5), (2, 64, 1536), (2, 4, 4, 1536)
 ```
 
 ## Citations

{alphagenome_pytorch-0.0.8 → alphagenome_pytorch-0.0.10}/README.md

@@ -14,13 +14,13 @@ $ pip install alphagenome-pytorch
 
 ```python
 import torch
-from alphagenome_pytorch import TransformerTower
+from alphagenome_pytorch import AlphaGenome
 
-transformer = TransformerTower(dim = 768, dim_pairwise = 128)
+model = AlphaGenome()
 
-single = torch.randn(2, 512, 768)
+dna = torch.randint(0, 5, (2, 8192))
 
-attended_single, attended_pairwise = transformer(single)
+pred_nucleotide, single, pairwise = model(dna) # (2, 8192, 5), (2, 64, 1536), (2, 4, 4, 1536)
 ```
 
 ## Citations

alphagenome_pytorch-0.0.10/alphagenome_pytorch/__init__.py

@@ -0,0 +1,11 @@
+from alphagenome_pytorch.alphagenome import (
+    AlphaGenome,
+    Attention,
+    PairwiseRowAttention,
+    RelativePosFeatures,
+    RotaryEmbedding,
+    FeedForward,
+    TransformerTower,
+    UpresBlock,
+    DownresBlock,
+)

{alphagenome_pytorch-0.0.8 → alphagenome_pytorch-0.0.10}/alphagenome_pytorch/alphagenome.py

@@ -4,7 +4,9 @@ from functools import partial
 import torch
 from torch import nn, cat, stack, arange, logspace
 import torch.nn.functional as F
-from torch.nn import Linear, Sequential, Module, ModuleList
+from torch.nn import Conv1d, Linear, Sequential, Module, ModuleList
+
+from torch.nn.utils.parametrize import register_parametrization
 
 from einx import add, multiply, greater
 from einops.layers.torch import Rearrange, Reduce
@@ -30,6 +32,9 @@ def exists(v):
 def divisible_by(num, den):
     return (num % den) == 0
 
+def last(arr):
+    return arr[-1]
+
 def is_odd(num):
     return not divisible_by(num, 2)
 
@@ -42,6 +47,104 @@ def default(v, d):
 def softclamp(t, value = 5.):
     return (t / value).tanh() * value
 
+# convolutional unet related
+
+class WeightStandardConv(Conv1d):
+    def __init__(
+        self,
+        dim,
+        dim_out,
+        width,
+        *args,
+        **kwargs
+    ):
+        super().__init__(dim, dim_out, width, *args, **kwargs)
+
+        register_parametrization(self, 'weight', nn.LayerNorm(self.weight.shape, elementwise_affine = False))
+
+class ConvBlock(Module):
+    def __init__(
+        self,
+        dim,
+        width = 5,
+        dim_out = None
+    ):
+        super().__init__()
+        assert is_odd(width)
+        dim_out = default(dim_out, dim)
+
+        conv_klass = Conv1d if width == 1 else WeightStandardConv
+
+        self.conv = conv_klass(dim, dim_out, width, padding = width // 2)
+
+    def forward(self, x):
+
+        x = F.gelu(x)
+        out = self.conv(x)
+        return out
+
+class DownresBlock(Module):
+    def __init__(
+        self,
+        dim,
+        channels_to_add = 128 # this is new as well? instead of doubling channels, they add 128 at a time, and use padding or slicing for the residual
+    ):
+        super().__init__()
+
+        dim_out = dim + channels_to_add
+        self.pad = channels_to_add
+
+        self.conv = ConvBlock(dim, width = 1, dim_out = dim_out)
+        self.conv_out = ConvBlock(dim_out, width = 1)
+
+        self.max_pool = Reduce('b d (n pool) -> b d n', 'max', pool = 2)
+
+    def forward(self, x):
+
+        residual = F.pad(x, (0, 0, 0, self.pad), value = 0.)
+
+        out = self.conv(x) + residual
+
+        out = self.conv_out(out) + out
+
+        return self.max_pool(out)
+
+class UpresBlock(Module):
+    def __init__(
+        self,
+        dim,
+        channels_to_remove = 128,
+        residual_scale_init = .9
+    ):
+        super().__init__()
+
+        dim_out = dim - channels_to_remove
+        self.pad = channels_to_remove
+
+        self.conv = ConvBlock(dim, width = 1, dim_out = dim_out)
+        self.unet_conv = ConvBlock(dim_out, width = 1)
+
+        self.conv_out = ConvBlock(dim_out, width = 1)
+
+        self.residual_scale = nn.Parameter(torch.ones(1,) * residual_scale_init)
+
+    def forward(
+        self,
+        x,
+        skip = None
+    ):
+
+        residual = x[:, :-self.pad]
+        out = self.conv(x) + residual
+
+        if exists(skip):
+            out = repeat(out, 'b c n -> b c (n upsample)', upsample = 2) * self.residual_scale
+            out = out + self.unet_conv(skip)
+
+        return self.conv_out(out) + out
+
+# position related
+
 def relative_shift(t):
     *leading_dims, seq_len, dim = t.shape
     t = F.pad(t, (1, 0), value = 0.)
@@ -284,7 +387,7 @@ class SingleToPairwise(Module):
             rel_q = relative_shift(einsum(q + q_rel_bias, rel_pos_encoding, 'b n h d, p h d -> b h n p'))
             rel_k = relative_shift(einsum(k + k_rel_bias, rel_pos_encoding, 'b n h d, p h d -> b h n p'))
 
-            rel_sim = add('b h i j, b h j i -> b h i j', rel_q, rel_k) * 0.5
+            rel_sim = add('b h i j, b h j i -> b i j h', rel_q, rel_k) * 0.5
 
             sim = sim + rel_sim
 
@@ -360,7 +463,7 @@ class TransformerTower(Module):
         dropout = 0.,
         ff_expansion_factor = 2.,
         max_positions = 8192,
-        dim_pairwise = None,
+        dim_pairwise = 128,
         pairwise_every_num_single_blocks = 2,   # how often to do a pairwise block
         single_to_pairwise_heads = 32,          # they did 32
         pool_size = 16,
@@ -456,8 +559,10 @@ class DNAEmbed(Module):
         super().__init__()
         assert is_odd(width)
         self.dim_input = dim_input
-        self.conv = nn.Conv1d(dim_input, dim, width, padding = width // 2)
-        self.pointwise = nn.Conv1d(dim, dim, 1)
+        self.conv = Conv1d(dim_input, dim, width, padding = width // 2)
+        self.pointwise = Conv1d(dim, dim, 1)
+
+        self.pool = Reduce('b d (n pool) -> b d n', 'max', pool = 2)
 
     def forward(
         self,
@@ -468,38 +573,96 @@ class DNAEmbed(Module):
 
         out = self.conv(x)
         out = out + self.pointwise(out)
-        return rearrange(out, 'b d n -> b n d')
+        pooled = self.pool(out) # think they downsample for dna embed block
+
+        return pooled, x
 
 # classes
 
 class AlphaGenome(Module):
     def __init__(
         self,
-        dim = 768,
+        dims: tuple[int, ...] = (
+            768,
+            896,
+            1024,
+            1152,
+            1280,
+            1408,
+            1536
+        ),
         basepairs = 5,
         dna_embed_width = 15,
-        dim_pairwise = None,
         transformer_kwargs: dict = dict()
     ):
         super().__init__()
+
         assert is_odd(dna_embed_width)
 
-        self.to_dna_embed = DNAEmbed(dim, dim_input = basepairs, width = dna_embed_width)
+        assert len(dims) >= 2
+        first_dim, *_, last_dim = dims
+
+        self.dna_embed = DNAEmbed(first_dim, dim_input = basepairs, width = dna_embed_width)
 
-        self.transformer = Transformer(
-            dim = dim,
-            dim_pairwise = dim_pairwise,
+        dim_with_input = (basepairs, *dims)
+        dim_pairs = zip(dim_with_input[:-1], dim_with_input[1:])
+
+        downs = []
+        ups = []
+
+        for layer_num, (dim_in, dim_out) in enumerate(dim_pairs, start = 1):
+            is_first = layer_num == 1
+            channel_diff = dim_out - dim_in
+
+            assert channel_diff > 0
+
+            if not is_first:
+                down = DownresBlock(dim_in, channels_to_add = channel_diff)
+                downs.append(down)
+
+            up = UpresBlock(dim_out, channels_to_remove = channel_diff)
+            ups.insert(0, up)
+
+
+        self.downs = ModuleList(downs)
+        self.ups = ModuleList(ups)
+
+        self.transformer = TransformerTower(
+            dim = last_dim,
             **transformer_kwargs
         )
 
     def forward(
         self,
-        seq,
-        pairwise
+        seq # Int['b n']
     ):
 
-        dna_embed = self.to_dna_embed(seq)
+        skips = []
+
+        # embed with one hot and add skip
+
+        x, skip = self.dna_embed(seq)
+        skips.append(skip)
+
+        # downs
+
+        for down in self.downs:
+            skips.append(x)
+            x = down(x)
+
+        x = rearrange(x, 'b d n -> b n d')
+
+        # attention
+
+        single, pairwise = self.transformer(x)
+
+        # ups with skips from down
+
+        x = rearrange(x, 'b n d -> b d n')
+
+        for up in self.ups:
+            x = up(x, skip = skips.pop())
 
-        attended = self.transformer(dna_embed)
+        pred = rearrange(x, 'b l n -> b n l') # 1bp resolution
 
-        return attended
+        return pred, single, pairwise

{alphagenome_pytorch-0.0.8 → alphagenome_pytorch-0.0.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphagenome-pytorch"
-version = "0.0.8"
+version = "0.0.10"
 description = "AlphaGenome"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

alphagenome_pytorch-0.0.10/tests/test_alphagenome.py

@@ -0,0 +1,36 @@
+import pytest
+
+import torch
+from alphagenome_pytorch.alphagenome import TransformerTower
+
+def test_attention():
+
+    transformer = TransformerTower(dim = 768, dim_pairwise = 128)
+
+    single = torch.randn(2, 512, 768)
+
+    single_repr, pairwise_repr = transformer(single)
+
+    assert single_repr.shape == (2, 512, 768)
+    assert pairwise_repr.shape == (2, 512 // 16, 512 // 16, 128)
+
+def test_down_up():
+    from alphagenome_pytorch.alphagenome import DownresBlock, UpresBlock
+    down = DownresBlock(64)
+    up = UpresBlock(64 + 128)
+
+    x = torch.randn(1, 64, 8)
+    assert up(down(x), x).shape == x.shape
+
+def test_alphagenome():
+    from alphagenome_pytorch import AlphaGenome
+
+    model = AlphaGenome()
+
+    dna = torch.randint(0, 5, (2, 8192))
+
+    pred_nucleotide_logits, single, pairwise = model(dna)
+
+    pred = pred_nucleotide_logits.argmax(dim = -1)
+
+    assert pred.shape == dna.shape

alphagenome_pytorch-0.0.8/alphagenome_pytorch/__init__.py

@@ -1,5 +0,0 @@
-from alphagenome_pytorch.alphagenome import (
-    AlphaGenome,
-    Attention,
-    TransformerTower
-)

alphagenome_pytorch-0.0.8/tests/test_alphagenome.py

@@ -1,14 +0,0 @@
-import pytest
-import torch
-from alphagenome_pytorch.alphagenome import TransformerTower
-
-def test_attention():
-
-    transformer = TransformerTower(dim = 768, dim_pairwise = 128)
-
-    single = torch.randn(2, 512, 768)
-
-    single_repr, pairwise_repr = transformer(single)
-
-    assert single_repr.shape == (2, 512, 768)
-    assert pairwise_repr.shape == (2, 512 // 16, 512 // 16, 128)