Add AutoregressiveTransformer

generative/networks/nets/__init__.py

@@ -10,6 +10,7 @@
 # limitations under the License.
 
 from .autoencoderkl import AutoencoderKL
+from .autoregressive_transformer import AutoregressiveTransformer
 from .diffusion_model_unet import DiffusionModelUNet
 from .patchgan_discriminator import MultiScalePatchDiscriminator, PatchDiscriminator
 from .vqvae import VQVAE

generative/networks/nets/autoregressive_transformer.py

@@ -0,0 +1,46 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch.nn as nn
+from x_transformers import Decoder, TransformerWrapper
+
+__all__ = ["AutoregressiveTransformer"]
+
+
+class AutoregressiveTransformer(nn.Module):
+    """Autoregressive Transformer model.
+
+    Args:
+        num_tokens: Number of tokens in the vocabulary.
+        max_seq_len: Maximum sequence length.
+        attn_layers_dim: Dimensionality of the attention layers.
+        attn_layers_depth: Number of attention layers.
+        attn_layers_heads: Number of attention heads.
+    """
+
+    def __init__(
+        self, num_tokens: int, max_seq_len: int, attn_layers_dim: int, attn_layers_depth: int, attn_layers_heads: int
+    ) -> None:
+        super().__init__()
+        self.num_tokens = num_tokens
+        self.max_seq_len = max_seq_len
+        self.attn_layers_dim = attn_layers_dim
+        self.attn_layers_depth = attn_layers_depth
+        self.attn_layers_heads = attn_layers_heads
+
+        self.model = TransformerWrapper(
+            num_tokens=self.num_tokens,
+            max_seq_len=self.max_seq_len,
+            attn_layers=Decoder(dim=self.attn_layers_dim, depth=self.attn_layers_depth, heads=self.attn_layers_heads),
+        )
+
+    def forward(self, x):
+        return self.model(x)

requirements-dev.txt

@@ -54,3 +54,4 @@ nni
 optuna
 git+https://github.com/Project-MONAI/MetricsReloaded@monai-support#egg=MetricsReloaded
 lpips==0.1.4
+x-transformers==1.8.1

tests/min_tests.py

@@ -33,6 +33,7 @@ def run_testsuit():
         "test_integration_workflows_adversarial",
         "test_latent_diffusion_inferer",
         "test_diffusion_inferer",
+        "test_autoregressive_transformer",
     ]
     assert sorted(exclude_cases) == sorted(set(exclude_cases)), f"Duplicated items in {exclude_cases}"
 

tests/test_autoregressive_transformer.py

@@ -0,0 +1,30 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import torch
+from monai.networks import eval_mode
+
+from generative.networks.nets import AutoregressiveTransformer
+
+
+class TestAutoregressiveTransformer(unittest.TestCase):
+    def test_shape_unconditioned_models(self):
+        net = AutoregressiveTransformer(
+            num_tokens=10, max_seq_len=16, attn_layers_dim=8, attn_layers_depth=2, attn_layers_heads=2
+        )
+        with eval_mode(net):
+            net.forward(torch.randint(0, 10, (1, 16)))
+
+
+if __name__ == "__main__":
+    unittest.main()