[Doc, BugFix] Fix tensordictmodule tutorial (#819)

Author: vmoens

docs/source/conf.py

@@ -182,5 +182,5 @@
 
 generate_knowledge_base_references("../../knowledge_base")
 generate_tutorial_references("../../tutorials/sphinx-tutorials/", "tutorial")
-generate_tutorial_references("../../tutorials/src/", "src")
+# generate_tutorial_references("../../tutorials/src/", "src")
 generate_tutorial_references("../../tutorials/media/", "media")

docs/source/content_generation.py

@@ -6,7 +6,7 @@
 FILE_DIR = os.path.dirname(__file__)
 KNOWLEDGE_GEN_DIR = "reference/generated/knowledge_base"
 TUTORIALS_GEN_DIR = "reference/generated/tutorials"
-TUTORIALS_SRC_GEN_DIR = "reference/generated/tutorials/src"
+# TUTORIALS_SRC_GEN_DIR = "reference/generated/tutorials/src"
 TUTORIALS_MEDIA_GEN_DIR = "reference/generated/tutorials/media"
 
 
@@ -71,8 +71,8 @@ def generate_tutorial_references(tutorial_path: str, file_type: str) -> None:
     # Create target dir
     if file_type == "tutorial":
         target_path = os.path.join(FILE_DIR, TUTORIALS_GEN_DIR)
-    elif file_type == "src":
-        target_path = os.path.join(FILE_DIR, TUTORIALS_SRC_GEN_DIR)
+    # elif file_type == "src":
+    #     target_path = os.path.join(FILE_DIR, TUTORIALS_SRC_GEN_DIR)
     else:
         target_path = os.path.join(FILE_DIR, TUTORIALS_MEDIA_GEN_DIR)
     Path(target_path).mkdir(parents=True, exist_ok=True)

tutorials/sphinx-tutorials/tensordict_module.py

@@ -333,16 +333,171 @@ def forward(self, x):
 #
 # We have let the positional encoders aside for simplicity.
 #
-# Let's first import the classical transformers blocks
-# (see ``src/transformer.py`` for more details.)
-
-from tutorials.src.transformer import (
-    Attention,
-    FFN,
-    SkipLayerNorm,
-    SplitHeads,
-    TokensToQKV,
-)
+# Let's re-write the classical transformers blocks:
+
+
+class TokensToQKV(nn.Module):
+    def __init__(self, to_dim, from_dim, latent_dim):
+        super().__init__()
+        self.q = nn.Linear(to_dim, latent_dim)
+        self.k = nn.Linear(from_dim, latent_dim)
+        self.v = nn.Linear(from_dim, latent_dim)
+
+    def forward(self, X_to, X_from):
+        Q = self.q(X_to)
+        K = self.k(X_from)
+        V = self.v(X_from)
+        return Q, K, V
+
+
+class SplitHeads(nn.Module):
+    def __init__(self, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+
+    def forward(self, Q, K, V):
+        batch_size, to_num, latent_dim = Q.shape
+        _, from_num, _ = K.shape
+        d_tensor = latent_dim // self.num_heads
+        Q = Q.reshape(batch_size, to_num, self.num_heads, d_tensor).transpose(1, 2)
+        K = K.reshape(batch_size, from_num, self.num_heads, d_tensor).transpose(1, 2)
+        V = V.reshape(batch_size, from_num, self.num_heads, d_tensor).transpose(1, 2)
+        return Q, K, V
+
+
+class Attention(nn.Module):
+    def __init__(self, latent_dim, to_dim):
+        super().__init__()
+        self.softmax = nn.Softmax(dim=-1)
+        self.out = nn.Linear(latent_dim, to_dim)
+
+    def forward(self, Q, K, V):
+        batch_size, n_heads, to_num, d_in = Q.shape
+        attn = self.softmax(Q @ K.transpose(2, 3) / d_in)
+        out = attn @ V
+        out = self.out(out.transpose(1, 2).reshape(batch_size, to_num, n_heads * d_in))
+        return out, attn
+
+
+class SkipLayerNorm(nn.Module):
+    def __init__(self, to_len, to_dim):
+        super().__init__()
+        self.layer_norm = nn.LayerNorm((to_len, to_dim))
+
+    def forward(self, x_0, x_1):
+        return self.layer_norm(x_0 + x_1)
+
+
+class FFN(nn.Module):
+    def __init__(self, to_dim, hidden_dim, dropout_rate=0.2):
+        super().__init__()
+        self.FFN = nn.Sequential(
+            nn.Linear(to_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, to_dim),
+            nn.Dropout(dropout_rate),
+        )
+
+    def forward(self, X):
+        return self.FFN(X)
+
+
+class AttentionBlock(nn.Module):
+    def __init__(self, to_dim, to_len, from_dim, latent_dim, num_heads):
+        super().__init__()
+        self.tokens_to_qkv = TokensToQKV(to_dim, from_dim, latent_dim)
+        self.split_heads = SplitHeads(num_heads)
+        self.attention = Attention(latent_dim, to_dim)
+        self.skip = SkipLayerNorm(to_len, to_dim)
+
+    def forward(self, X_to, X_from):
+        Q, K, V = self.tokens_to_qkv(X_to, X_from)
+        Q, K, V = self.split_heads(Q, K, V)
+        out, attention = self.attention(Q, K, V)
+        out = self.skip(X_to, out)
+        return out
+
+
+class EncoderTransformerBlock(nn.Module):
+    def __init__(self, to_dim, to_len, latent_dim, num_heads):
+        super().__init__()
+        self.attention_block = AttentionBlock(
+            to_dim, to_len, to_dim, latent_dim, num_heads
+        )
+        self.FFN = FFN(to_dim, 4 * to_dim)
+        self.skip = SkipLayerNorm(to_len, to_dim)
+
+    def forward(self, X_to):
+        X_to = self.attention_block(X_to, X_to)
+        X_out = self.FFN(X_to)
+        return self.skip(X_out, X_to)
+
+
+class DecoderTransformerBlock(nn.Module):
+    def __init__(self, to_dim, to_len, from_dim, latent_dim, num_heads):
+        super().__init__()
+        self.attention_block = AttentionBlock(
+            to_dim, to_len, from_dim, latent_dim, num_heads
+        )
+        self.encoder_block = EncoderTransformerBlock(
+            to_dim, to_len, latent_dim, num_heads
+        )
+
+    def forward(self, X_to, X_from):
+        X_to = self.attention_block(X_to, X_from)
+        X_to = self.encoder_block(X_to)
+        return X_to
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(self, num_blocks, to_dim, to_len, latent_dim, num_heads):
+        super().__init__()
+        self.encoder = nn.ModuleList(
+            [
+                EncoderTransformerBlock(to_dim, to_len, latent_dim, num_heads)
+                for i in range(num_blocks)
+            ]
+        )
+
+    def forward(self, X_to):
+        for i in range(len(self.encoder)):
+            X_to = self.encoder[i](X_to)
+        return X_to
+
+
+class TransformerDecoder(nn.Module):
+    def __init__(self, num_blocks, to_dim, to_len, from_dim, latent_dim, num_heads):
+        super().__init__()
+        self.decoder = nn.ModuleList(
+            [
+                DecoderTransformerBlock(to_dim, to_len, from_dim, latent_dim, num_heads)
+                for i in range(num_blocks)
+            ]
+        )
+
+    def forward(self, X_to, X_from):
+        for i in range(len(self.decoder)):
+            X_to = self.decoder[i](X_to, X_from)
+        return X_to
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self, num_blocks, to_dim, to_len, from_dim, from_len, latent_dim, num_heads
+    ):
+        super().__init__()
+        self.encoder = TransformerEncoder(
+            num_blocks, to_dim, to_len, latent_dim, num_heads
+        )
+        self.decoder = TransformerDecoder(
+            num_blocks, from_dim, from_len, to_dim, latent_dim, num_heads
+        )
+
+    def forward(self, X_to, X_from):
+        X_to = self.encoder(X_to)
+        X_out = self.decoder(X_from, X_to)
+        return X_out
+
 
 ###############################################################################
 # We first create the ``AttentionBlockTensorDict``, the attention block using
@@ -608,8 +763,6 @@ def __init__(
 # Benchmarking
 # ------------------------------
 
-from tutorials.src.transformer import Transformer
-
 ###############################################################################
 
 to_dim = 5