Max margin training

yzhangcs · yzhangcs · commit c6ca288c4da9 · 2021-07-19T22:38:57.000+08:00
diff --git a/supar/cmds/crf_con.py b/supar/cmds/crf_con.py
@@ -26,6 +26,7 @@ def main():
     subparser.add_argument('--unk', default='unk', help='unk token in pretrained embeddings')
     subparser.add_argument('--n-embed', default=100, type=int, help='dimension of embeddings')
     subparser.add_argument('--bert', default='bert-base-cased', help='which BERT model to use')
+    subparser.add_argument('--loss', choices=['crf', 'max-margin'], default='crf', help='loss for global training')
     # evaluate
     subparser = subparsers.add_parser('evaluate', help='Evaluate the specified parser and dataset.')
     subparser.add_argument('--buckets', default=8, type=int, help='max num of buckets to use')
diff --git a/supar/models/con.py b/supar/models/con.py
@@ -175,7 +175,7 @@ def loss(self, s_span, s_label, charts, mask, mbr=True):
             s_span (~torch.Tensor): ``[batch_size, seq_len, seq_len]``.
                 Scores of all constituents.
             s_label (~torch.Tensor): ``[batch_size, seq_len, seq_len, n_labels]``.
-                Scores of all labels on each constituent.
+                Scores of all constituent labels.
             charts (~torch.LongTensor): ``[batch_size, seq_len, seq_len]``.
                 The tensor of gold-standard labels. Positions without labels are filled with -1.
             mask (~torch.BoolTensor): ``[batch_size, seq_len, seq_len]``.
@@ -190,8 +190,13 @@ def loss(self, s_span, s_label, charts, mask, mbr=True):
         """
 
         span_mask = charts.ge(0) & mask
-        span_dist = CRFConstituency(s_span, mask)
-        span_loss = -span_dist.log_prob(span_mask).sum() / mask[:, 0].sum()
+        if self.args.loss == 'crf':
+            span_dist = CRFConstituency(s_span, mask)
+            span_loss = -span_dist.log_prob(span_mask).sum()
+        elif self.args.loss == 'max-margin':
+            span_dist = CRFConstituency(s_span + torch.full_like(s_span, 1) - span_mask.float(), mask)
+            span_loss = span_dist.max.sum() - s_span[span_mask].sum()
+        span_loss = span_loss / mask[:, 0].sum()
         span_probs = span_dist.marginals if mbr else s_span
         label_loss = self.criterion(s_label[span_mask], charts[span_mask])
         loss = span_loss + label_loss
@@ -204,7 +209,7 @@ def decode(self, s_span, s_label, mask):
             s_span (~torch.Tensor): ``[batch_size, seq_len, seq_len]``.
                 Scores of all constituents.
             s_label (~torch.Tensor): ``[batch_size, seq_len, seq_len, n_labels]``.
-                Scores of all labels on each constituent.
+                Scores of all constituent labels.
             mask (~torch.BoolTensor): ``[batch_size, seq_len, seq_len]``.
                 The mask for covering the unpadded tokens in each chart.
 
@@ -406,7 +411,7 @@ def loss(self, s_span, s_pair, s_label, charts, mask):
             s_pair (~torch.Tensor): ``[batch_size, seq_len, seq_len, seq_len]``.
                 Scores of second-order triples.
             s_label (~torch.Tensor): ``[batch_size, seq_len, seq_len, n_labels]``.
-                Scores of all labels on each constituent.
+                Scores of all constituent labels.
             charts (~torch.LongTensor): ``[batch_size, seq_len, seq_len]``.
                 The tensor of gold-standard labels. Positions without labels are filled with -1.
             mask (~torch.BoolTensor): ``[batch_size, seq_len, seq_len]``.
@@ -430,7 +435,7 @@ def decode(self, s_span, s_label, mask):
             s_span (~torch.Tensor): ``[batch_size, seq_len, seq_len]``.
                 Scores of all constituents.
             s_label (~torch.Tensor): ``[batch_size, seq_len, seq_len, n_labels]``.
-                Scores of all labels on each constituent.
+                Scores of all constituent labels.
             mask (~torch.BoolTensor): ``[batch_size, seq_len, seq_len]``.
                 The mask for covering the unpadded tokens in each chart.
 
diff --git a/supar/structs/crf.py b/supar/structs/crf.py
@@ -289,10 +289,10 @@ def __repr__(self):
 
     @lazy_property
     def argmax(self):
-        return [sorted(i.nonzero().tolist(), key=lambda x:(x[0], -x[1])) for i in self.backward(self.max.sum())]
+        return [sorted(torch.nonzero(i).tolist(), key=lambda x:(x[0], -x[1])) for i in self.backward(self.max.sum())]
 
     def topk(self, k):
-        return list(zip(*[[sorted(i.nonzero().tolist(), key=lambda x:(x[0], -x[1])) for i in self.backward(i)]
+        return list(zip(*[[sorted(torch.nonzero(i).tolist(), key=lambda x:(x[0], -x[1])) for i in self.backward(i)]
                           for i in self.kmax(k).sum(0)]))
 
     def score(self, value):
