Update real_time_encoder_transformer.py

Author: ajatshatru01

neural_network/real_time_encoder_transformer.py

@@ -1,4 +1,3 @@
-# imports
 import math
 
 import torch
@@ -16,7 +15,6 @@ class Time2Vec(nn.Module):
     >>> output.shape
     torch.Size([1, 3, 4])
     """
-
     def __init__(self, d_model: int) -> None:
         super().__init__()
         self.w0 = nn.Parameter(torch.randn(1, 1))
@@ -41,7 +39,6 @@ class PositionwiseFeedForward(nn.Module):
     >>> out.shape
     torch.Size([4, 10, 8])
     """
-
     def __init__(self, d_model: int, hidden: int, drop_prob: float = 0.1) -> None:
         super().__init__()
         self.fc1 = nn.Linear(d_model, hidden)
@@ -62,29 +59,32 @@ class ScaleDotProductAttention(nn.Module):
 
     >>> import torch
     >>> attn = ScaleDotProductAttention()
-    >>> q = torch.rand(2, 8, 10, 16)
-    >>> k = torch.rand(2, 8, 10, 16)
-    >>> v = torch.rand(2, 8, 10, 16)
-    >>> ctx, attn_w = attn.forward(q, k, v)
+    >>> query_tensor = torch.rand(2, 8, 10, 16)
+    >>> key_tensor = torch.rand(2, 8, 10, 16)
+    >>> value_tensor = torch.rand(2, 8, 10, 16)
+    >>> ctx, attn_w = attn.forward(query_tensor, key_tensor, value_tensor)
     >>> ctx.shape
     torch.Size([2, 8, 10, 16])
     """
-
     def __init__(self) -> None:
         super().__init__()
         self.softmax = nn.Softmax(dim=-1)
 
     def forward(
-        self, q: Tensor, k: Tensor, v: Tensor, mask: Tensor = None
+        self,
+        query_tensor: Tensor,
+        key_tensor: Tensor,
+        value_tensor: Tensor,
+        mask: Tensor = None,
     ) -> tuple[Tensor, Tensor]:
-        _, _, _, d_k = k.size()
-        scores = (q @ k.transpose(2, 3)) / math.sqrt(d_k)
+        _, _, _, d_k = key_tensor.size()
+        scores = (query_tensor @ key_tensor.transpose(2, 3)) / math.sqrt(d_k)
 
         if mask is not None:
             scores = scores.masked_fill(mask == 0, -1e9)
 
         attn = self.softmax(scores)
-        context = attn @ v
+        context = attn @ value_tensor
         return context, attn
 
 
@@ -94,12 +94,11 @@ class MultiHeadAttention(nn.Module):
 
     >>> import torch
     >>> attn = MultiHeadAttention(16, 4)
-    >>> q = torch.rand(2, 10, 16)
-    >>> out = attn.forward(q, q, q)
+    >>> query_tensor = torch.rand(2, 10, 16)
+    >>> out = attn.forward(query_tensor, query_tensor, query_tensor)
     >>> out.shape
     torch.Size([2, 10, 16])
     """
-
     def __init__(self, d_model: int, n_head: int) -> None:
         super().__init__()
         self.n_head = n_head
@@ -109,22 +108,34 @@ def __init__(self, d_model: int, n_head: int) -> None:
         self.w_v = nn.Linear(d_model, d_model)
         self.w_out = nn.Linear(d_model, d_model)
 
-    def forward(self, q: Tensor, k: Tensor, v: Tensor, mask: Tensor = None) -> Tensor:
-        q, k, v = self.w_q(q), self.w_k(k), self.w_v(v)
-        q, k, v = self.split_heads(q), self.split_heads(k), self.split_heads(v)
+    def forward(
+        self,
+        query_tensor: Tensor,
+        key_tensor: Tensor,
+        value_tensor: Tensor,
+        mask: Tensor = None,
+    ) -> Tensor:
+        query_tensor, key_tensor, value_tensor = (
+            self.w_q(query_tensor),
+            self.w_k(key_tensor),
+            self.w_v(value_tensor),
+        )
+        query_tensor = self.split_heads(query_tensor)
+        key_tensor = self.split_heads(key_tensor)
+        value_tensor = self.split_heads(value_tensor)
 
-        context, _ = self.attn(q, k, v, mask)
+        context, _ = self.attn(query_tensor, key_tensor, value_tensor, mask)
         out = self.w_out(self.concat_heads(context))
         return out
 
-    def split_heads(self, x: Tensor) -> Tensor:
-        batch, seq_len, d_model = x.size()
+    def split_heads(self, input_tensor: Tensor) -> Tensor:
+        batch, seq_len, d_model = input_tensor.size()
         d_k = d_model // self.n_head
-        return x.view(batch, seq_len, self.n_head, d_k).transpose(1, 2)
+        return input_tensor.view(batch, seq_len, self.n_head, d_k).transpose(1, 2)
 
-    def concat_heads(self, x: Tensor) -> Tensor:
-        batch, n_head, seq_len, d_k = x.size()
-        return x.transpose(1, 2).contiguous().view(batch, seq_len, n_head * d_k)
+    def concat_heads(self, input_tensor: Tensor) -> Tensor:
+        batch, n_head, seq_len, d_k = input_tensor.size()
+        return input_tensor.transpose(1, 2).contiguous().view(batch, seq_len, n_head * d_k)
 
 
 class LayerNorm(nn.Module):
@@ -138,7 +149,6 @@ class LayerNorm(nn.Module):
     >>> out.shape
     torch.Size([4, 10, 8])
     """
-
     def __init__(self, d_model: int, eps: float = 1e-12) -> None:
         super().__init__()
         self.gamma = nn.Parameter(torch.ones(d_model))
@@ -148,9 +158,7 @@ def __init__(self, d_model: int, eps: float = 1e-12) -> None:
     def forward(self, input_tensor: Tensor) -> Tensor:
         mean = input_tensor.mean(-1, keepdim=True)
         var = input_tensor.var(-1, unbiased=False, keepdim=True)
-        return (
-            self.gamma * (input_tensor - mean) / torch.sqrt(var + self.eps) + self.beta
-        )
+        return self.gamma * (input_tensor - mean) / torch.sqrt(var + self.eps) + self.beta
 
 
 class TransformerEncoderLayer(nn.Module):
@@ -164,7 +172,6 @@ class TransformerEncoderLayer(nn.Module):
     >>> out.shape
     torch.Size([4, 10, 8])
     """
-
     def __init__(
         self,
         d_model: int,
@@ -198,7 +205,6 @@ class TransformerEncoder(nn.Module):
     >>> out.shape
     torch.Size([4, 10, 8])
     """
-
     def __init__(
         self,
         d_model: int,
@@ -235,14 +241,11 @@ class AttentionPooling(nn.Module):
     >>> weights.shape
     torch.Size([4, 10])
     """
-
     def __init__(self, d_model: int) -> None:
         super().__init__()
         self.attn_score = nn.Linear(d_model, 1)
 
-    def forward(
-        self, input_tensor: Tensor, mask: Tensor = None
-    ) -> tuple[Tensor, Tensor]:
+    def forward(self, input_tensor: Tensor, mask: Tensor = None) -> tuple[Tensor, Tensor]:
         attn_weights = torch.softmax(self.attn_score(input_tensor).squeeze(-1), dim=-1)
 
         if mask is not None:
@@ -264,7 +267,6 @@ class EEGTransformer(nn.Module):
     >>> out.shape
     torch.Size([2, 1])
     """
-
     def __init__(
         self,
         feature_dim: int,
@@ -286,16 +288,9 @@ def __init__(
         self.pooling = AttentionPooling(d_model)
         self.output_layer = nn.Linear(d_model, output_dim)
 
-    def forward(
-        self, input_tensor: Tensor, mask: Tensor = None
-    ) -> tuple[Tensor, Tensor]:
+    def forward(self, input_tensor: Tensor, mask: Tensor = None) -> tuple[Tensor, Tensor]:
         b, t, _ = input_tensor.size()
-        t_idx = (
-            torch.arange(t, device=input_tensor.device)
-            .view(1, t, 1)
-            .expand(b, t, 1)
-            .float()
-        )
+        t_idx = torch.arange(t, device=input_tensor.device).view(1, t, 1).expand(b, t, 1).float()
         time_emb = self.time2vec(t_idx)
         x = self.input_proj(input_tensor) + time_emb
         x = self.encoder(x, mask)