update mask operation

yusugomori · Mar 18, 2019 · 74fe26c · 74fe26c
1 parent cc48692
commit 74fe26c
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 7 deletions.
diff --git a/models/encoder_decoder_attention.py b/models/encoder_decoder_attention.py
@@ -16,9 +16,9 @@ def __init__(self,
                  input_dim,
                  hidden_dim,
                  output_dim,
-                 device='cpu',
                  bos_value=1,
-                 max_len=20):
+                 max_len=20,
+                 device='cpu'):
         super().__init__()
         self.device = device
         self.encoder = Encoder(input_dim, hidden_dim, device=device)

diff --git a/models/encoder_decoder_lstm.py b/models/encoder_decoder_lstm.py
@@ -15,7 +15,6 @@ def __init__(self,
                  input_dim,
                  hidden_dim,
                  output_dim,
-                 device,
                  bos_value=1,
                  max_len=20,
                  device='cpu'):

diff --git a/models/layers/Attention.py b/models/layers/Attention.py
@@ -38,8 +38,10 @@ def forward(self, ht, hs, source=None, pad_value=0):
 
         score = torch.exp(score)
         if source is not None:
-            mask_source = (source.t() != pad_value).unsqueeze(0)
-            score = score * mask_source.float().to(self.device)
+            # mask_source = (source.t() != pad_value).unsqueeze(0)
+            # score = score * mask_source.float().to(self.device)
+            mask_source = source.t().eq(pad_value).unsqueeze(0)
+            score.data.masked_fill_(mask_source, 0)
 
         a = score / torch.sum(score, dim=-1, keepdim=True)
         c = torch.einsum('jik,kil->jil', (a, hs))

diff --git a/models/layers/DotProductAttention.py b/models/layers/DotProductAttention.py
@@ -21,7 +21,8 @@ def forward(self, q, k, v, mask=None):
             # in source-target-attention, source is `k` and `v`
             if len(mask.size()) == 2:
                 mask = mask.unsqueeze(0)
-            score = score * mask.type(torch.Tensor).to(self.device)
+            # score = score * mask.float().to(self.device)
+            score.data.masked_fill_(mask, 0)
 
         a = score / torch.sum(score, dim=-1, keepdim=True)
         c = torch.einsum('jik,kil->jil', (a, v))

diff --git a/models/layers/ScaledDotProductAttention.py b/models/layers/ScaledDotProductAttention.py
@@ -23,7 +23,8 @@ def forward(self, q, k, v, mask=None):
             # in source-target-attention, source is `k` and `v`
             if len(mask.size()) == 2:
                 mask = mask.unsqueeze(0)
-            score = score * mask.type(torch.Tensor).to(self.device)
+            # score = score * mask.float().to(self.device)
+            score.data.masked_fill_(mask, 0)
 
         a = score / torch.sum(score, dim=-1, keepdim=True)
         c = torch.einsum('jik,kil->jil', (a, v))