modify transpose

Author: serotoninpm

.gitignore

@@ -4,3 +4,4 @@ venv/
 .data/
 *.pt
 __pycache__
+result/

models/layers/multi_head_attention.py

@@ -48,7 +48,7 @@ def split(self, tensor):
         batch_size, length, d_model = tensor.size()
 
         d_tensor = d_model // self.n_head
-        tensor = tensor.view(batch_size, self.n_head, length, d_tensor)
+        tensor = tensor.view(batch_size, length, self.n_head, d_tensor).transpose(1, 2)
         # it is similar with group convolution (split by number of heads)
 
         return tensor
@@ -63,5 +63,5 @@ def concat(self, tensor):
         batch_size, head, length, d_tensor = tensor.size()
         d_model = head * d_tensor
 
-        tensor = tensor.view(batch_size, length, d_model)
+        tensor = tensor.transpose(1, 2).contiguous().view(batch_size, length, d_model)
         return tensor

models/layers/scale_dot_product_attention.py

@@ -19,15 +19,15 @@ class ScaleDotProductAttention(nn.Module):
 
     def __init__(self):
         super(ScaleDotProductAttention, self).__init__()
-        self.softmax = nn.Softmax()
+        self.softmax = nn.Softmax(dim=-1)
 
     def forward(self, q, k, v, mask=None, e=1e-12):
         # input is 4 dimension tensor
         # [batch_size, head, length, d_tensor]
         batch_size, head, length, d_tensor = k.size()
 
         # 1. dot product Query with Key^T to compute similarity
-        k_t = k.view(batch_size, head, d_tensor, length)  # transpose
+        k_t = k.transpose(2, 3)  # transpose
         score = (q @ k_t) / math.sqrt(d_tensor)  # scaled dot product
 
         # 2. apply masking (opt)