More efficient biaffine operation via einsum

zysite · zysite · commit c22c4000b2c7 · 2019-12-11T04:41:20.000Z
diff --git a/parser/modules/biaffine.py b/parser/modules/biaffine.py
@@ -32,15 +32,11 @@ def reset_parameters(self):
 
     def forward(self, x, y):
         if self.bias_x:
-            x = torch.cat([x, x.new_ones(x.shape[:-1]).unsqueeze(-1)], -1)
+            x = torch.cat((x, torch.ones_like(x[..., :1])), -1)
         if self.bias_y:
-            y = torch.cat([y, y.new_ones(y.shape[:-1]).unsqueeze(-1)], -1)
-        # [batch_size, 1, seq_len, d]
-        x = x.unsqueeze(1)
-        # [batch_size, 1, seq_len, d]
-        y = y.unsqueeze(1)
+            y = torch.cat((y, torch.ones_like(y[..., :1])), -1)
         # [batch_size, n_out, seq_len, seq_len]
-        s = x @ self.weight @ y.transpose(-1, -2)
+        s = torch.einsum('bxi,oij,byj->boxy', x, self.weight, y)
         # remove dim 1 if n_out == 1
         s = s.squeeze(1)
 
