Optimize einsum

yzhangcs · yzhangcs · commit a8e6f443febf · 2021-05-14T17:12:59.000+08:00
diff --git a/setup.py b/setup.py
@@ -27,6 +27,7 @@
         'transformers>=4.0.0',
         'nltk',
         'stanza',
+        'opt_einsum',
         'dill'],
     entry_points={
         'console_scripts': [
diff --git a/supar/modules/affine.py b/supar/modules/affine.py
@@ -2,6 +2,7 @@
 
 import torch
 import torch.nn as nn
+from opt_einsum import contract
 
 
 class Biaffine(nn.Module):
@@ -71,7 +72,7 @@ def forward(self, x, y):
         if self.bias_y:
             y = torch.cat((y, torch.ones_like(y[..., :1])), -1)
         # [batch_size, n_out, seq_len, seq_len]
-        s = torch.einsum('bxi,oij,byj->boxy', x, self.weight, y) / self.n_in ** self.scale
+        s = contract('bxi,oij,byj->boxy', x, self.weight, y) / self.n_in ** self.scale
         # remove dim 1 if n_out == 1
         s = s.squeeze(1)
 
@@ -145,9 +146,9 @@ def forward(self, x, y, z):
             x = torch.cat((x, torch.ones_like(x[..., :1])), -1)
         if self.bias_y:
             y = torch.cat((y, torch.ones_like(y[..., :1])), -1)
-        w = torch.einsum('bzk,oikj->bozij', z, self.weight)
+        w = contract('bzk,oikj->bozij', z, self.weight)
         # [batch_size, n_out, seq_len, seq_len, seq_len]
-        s = torch.einsum('bxi,bozij,byj->bozxy', x, w, y) / self.n_in ** self.scale
+        s = contract('bxi,bozij,byj->bozxy', x, w, y) / self.n_in ** self.scale
         # remove dim 1 if n_out == 1
         s = s.squeeze(1)
 
