fix pcnn bug

Author: Tianyu Gao

.gitignore

@@ -116,6 +116,7 @@ ckpt
 benchmark/tacred
 *.swp
 
+<<<<<<< HEAD
 # data and pretrain
 pretrain
 benchmark
@@ -127,3 +128,8 @@ benchmark
 
 # package
 opennre-egg.info
+=======
+# debug
+benchmark/nyt10-ori
+train_nyt10_pcnn_att_ori.py
+>>>>>>> ebf8370... fix pcnn bug

opennre/encoder/cnn_encoder.py

@@ -18,7 +18,8 @@ def __init__(self,
                  word2vec=None,
                  kernel_size=3, 
                  padding_size=1,
-                 dropout=0.5):
+                 dropout=0,
+                 activation_function=F.relu):
         """
         Args:
             token2id: dictionary of token->idx mapping
@@ -33,12 +34,13 @@ def __init__(self,
         """
         # Hyperparameters
         super(CNNEncoder, self).__init__(token2id, max_length, hidden_size, word_size, position_size, blank_padding, word2vec)
-        self.dropout = dropout
+        self.drop = nn.Dropout(dropout)
         self.kernel_size = kernel_size
         self.padding_size = padding_size
+        self.act = activation_function
 
-        self.conv = CNN(self.input_size, self.hidden_size, self.dropout, self.kernel_size, self.padding_size, activation_function = F.relu)
-        self.pool = MaxPool(self.max_length)
+        self.conv = nn.Conv1d(self.input_size, self.hidden_size, self.kernel_size, padding=self.padding)
+        self.pool = nn.MaxPool1d(self.max_length)
 
     def forward(self, token, pos1, pos2):
         """
@@ -55,9 +57,11 @@ def forward(self, token, pos1, pos2):
         x = torch.cat([self.word_embedding(token), 
                        self.pos1_embedding(pos1), 
                        self.pos2_embedding(pos2)], 2) # (B, L, EMBED)
-        x = self.conv(x) # (B, L, EMBED)
-        x = self.pool(x) # (B, EMBED)
+        x = x.transpose(1, 2) # (B, EMBED, L)
+        x = self.act(self.conv(x)) # (B, H, L)
+        x = self.pool(x).squeeze(-1)
+        x = self.drop(x)
         return x
 
     def tokenize(self, item):
-        return super().tokenize(item)
+        return super().tokenize(item)

opennre/encoder/pcnn_encoder.py

@@ -18,7 +18,8 @@ def __init__(self,
                  word2vec=None,
                  kernel_size=3, 
                  padding_size=1,
-                 dropout=0.5):
+                 dropout=0.0,
+                 activation_function=F.relu):
         """
         Args:
             token2id: dictionary of token->idx mapping
@@ -33,12 +34,19 @@ def __init__(self,
         """
         # hyperparameters
         super().__init__(token2id, max_length, hidden_size, word_size, position_size, blank_padding, word2vec)
-        self.dropout = dropout
+        self.drop = nn.Dropout(dropout)
         self.kernel_size = kernel_size
         self.padding_size = padding_size
+        self.act = activation_function
+    
+        self.conv = nn.Conv1d(self.input_size, self.hidden_size, self.kernel_size, padding=self.padding_size)
+        self.pool = nn.MaxPool1d(self.max_length)
+        self.mask_embedding = nn.Embedding(4, 3)
+        self.mask_embedding.weight.data.copy_(torch.FloatTensor([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]]))
+        self.mask_embedding.weight.requires_grad = False
+        self._minus = -100
 
-        self.conv = CNN(self.input_size, self.hidden_size, self.dropout, self.kernel_size, self.padding_size, activation_function=F.relu)
-        self.pool = MaxPool(self.max_length, 3)
+        self.hidden_size *= 3
 
     def forward(self, token, pos1, pos2, mask):
         """
@@ -55,8 +63,17 @@ def forward(self, token, pos1, pos2, mask):
         x = torch.cat([self.word_embedding(token), 
                        self.pos1_embedding(pos1), 
                        self.pos2_embedding(pos2)], 2) # (B, L, EMBED)
-        x = self.conv(x) # (B, L, EMBED)
-        x = self.pool(x) # (B, EMBED)
+        x = x.transpose(1, 2) # (B, EMBED, L)
+        x = self.conv(x) # (B, H, L)
+
+        mask = 1 - self.mask_embedding(mask).transpose(1, 2) # (B, L) -> (B, L, 3) -> (B, 3, L)
+        pool1 = self.pool(self.act(x + self._minus * mask[:, 0:1, :])) # (B, H, 1)
+        pool2 = self.pool(self.act(x + self._minus * mask[:, 1:2, :]))
+        pool3 = self.pool(self.act(x + self._minus * mask[:, 2:3, :]))
+        x = torch.cat([pool1, pool2, pool3], 1) # (B, 3H, 1)
+        x = x.squeeze(2) # (B, 3H)
+        x = self.drop(x)
+
         return x
 
     def tokenize(self, item):
@@ -69,25 +86,72 @@ def tokenize(self, item):
         Return:
             Name of the relation of the sentence
         """
-        # Sentence -> token
-        indexed_tokens, pos1, pos2 = super().tokenize(item)
-        sentence = item['text']
+        if 'text' in item:
+            sentence = item['text']
+            is_token = False
+        else:
+            sentence = item['token']
+            is_token = True
         pos_head = item['h']['pos']
-        pos_tail = item['t']['pos']        
+        pos_tail = item['t']['pos']
 
-        # Mask
+        # Sentence -> token
+        if not is_token:
+            if pos_head[0] > pos_tail[0]:
+                pos_min, pos_max = [pos_tail, pos_head]
+                rev = True
+            else:
+                pos_min, pos_max = [pos_head, pos_tail]
+                rev = False
+            sent_0 = self.tokenizer.tokenize(sentence[:pos_min[0]])
+            sent_1 = self.tokenizer.tokenize(sentence[pos_min[1]:pos_max[0]])
+            sent_2 = self.tokenizer.tokenize(sentence[pos_max[1]:])
+            ent_0 = self.tokenizer.tokenize(sentence[pos_min[0]:pos_min[1]])
+            ent_1 = self.tokenizer.tokenize(sentence[pos_max[0]:pos_max[1]])
+            tokens = sent_0 + ent_0 + sent_1 + ent_1 + sent_2
+            if rev:
+                pos_tail = [len(sent_0), len(sent_0) + len(ent_0)]
+                pos_head = [len(sent_0) + len(ent_0) + len(sent_1), len(sent_0) + len(ent_0) + len(sent_1) + len(ent_1)]
+            else:
+                pos_head = [len(sent_0), len(sent_0) + len(ent_0)]
+                pos_tail = [len(sent_0) + len(ent_0) + len(sent_1), len(sent_0) + len(ent_0) + len(sent_1) + len(ent_1)]           
+        else:
+            tokens = sentence
+
+        # Token -> index
+        if self.blank_padding:
+            indexed_tokens = self.tokenizer.convert_tokens_to_ids(tokens, self.max_length, self.token2id['[PAD]'], self.token2id['[UNK]'])
+        else:
+            indexed_tokens = self.tokenizer.convert_tokens_to_ids(tokens, unk_id = self.token2id['[UNK]'])
+
+        # Position -> index
+        pos1 = []
+        pos2 = []
         pos1_in_index = min(pos_head[0], self.max_length)
         pos2_in_index = min(pos_tail[0], self.max_length)
+        for i in range(len(tokens)):
+            pos1.append(min(i - pos1_in_index + self.max_length, 2 * self.max_length - 1))
+            pos2.append(min(i - pos2_in_index + self.max_length, 2 * self.max_length - 1))
+
+        if self.blank_padding:                
+            while len(pos1) < self.max_length:
+                pos1.append(0)
+            while len(pos2) < self.max_length:
+                pos2.append(0)
+            indexed_tokens = indexed_tokens[:self.max_length]
+            pos1 = pos1[:self.max_length]
+            pos2 = pos2[:self.max_length]
 
+        indexed_tokens = torch.tensor(indexed_tokens).long().unsqueeze(0) # (1, L)
+        pos1 = torch.tensor(pos1).long().unsqueeze(0) # (1, L)
+        pos2 = torch.tensor(pos2).long().unsqueeze(0) # (1, L)
+        
+        # Mask
         mask = []
-        pos_min = min(pos1_in_index, pos2_in_index)
-        pos_max = max(pos1_in_index, pos2_in_index)
-        for i in range(len(indexed_tokens)):
-            if pos1[0][i] == 0:
-                break
-            if i <= pos_min:
+        for i in range(len(tokens)):
+            if i <= pos_min[0]:
                 mask.append(1)
-            elif i <= pos_max:
+            elif i <= pos_max[0]:
                 mask.append(2)
             else:
                 mask.append(3)

opennre/module/pool/max_pool.py

@@ -30,15 +30,24 @@ def forward(self, x, mask=None):
             output features: (B, H_EMBED)
         """
         # Check size of tensors
-        if mask == None or self.segment_num == None or self.segment_num == 1:
+        if mask is None or self.segment_num is None or self.segment_num == 1:
             x = x.transpose(1, 2) # (B, L, I_EMBED) -> (B, I_EMBED, L)
             x = self.pool(x).squeeze(-1) # (B, I_EMBED, 1) -> (B, I_EMBED)
             return x
         else:
-            B, L, I_EMBED = x.size()[:2]
-            mask = 1 - self.mask_embedding(mask).transpose(1, 2).unsqueeze(2) # (B, L) -> (B, L, S) -> (B, S, L) -> (B, S, 1, L)
-            x = x.transpose(1, 2).unsqueeze(1) # (B, L, I_EMBED) -> (B, I_EMBED, L) -> (B, 1, I_EMBED, L)
-            x = (x + self._minus * mask).view([-1, I_EMBED, L]) # (B, S, I_EMBED, L) -> (B * S, I_EMBED, L)
-            x = self.pool(x).squeeze(-1) # (B * S, I_EMBED, 1) -> (B * S, I_EMBED)
-            x = x.view([B, -1])  # (B, S * I_EMBED)
-            return x
+            B, L, I_EMBED = x.size()[:3]
+            # mask = 1 - self.mask_embedding(mask).transpose(1, 2).unsqueeze(2) # (B, L) -> (B, L, S) -> (B, S, L) -> (B, S, 1, L)
+            # x = x.transpose(1, 2).unsqueeze(1) # (B, L, I_EMBED) -> (B, I_EMBED, L) -> (B, 1, I_EMBED, L)
+            # x = (x + self._minus * mask).contiguous().view([-1, I_EMBED, L]) # (B, S, I_EMBED, L) -> (B * S, I_EMBED, L)
+            # x = self.pool(x).squeeze(-1) # (B * S, I_EMBED, 1) -> (B * S, I_EMBED)
+            # x = x.view([B, -1])  # (B, S * I_EMBED)
+            # return x
+            mask = 1 - self.mask_embedding(mask).transpose(1, 2)
+            x = x.transpose(1, 2)
+            pool1 = self.pool(x + self._minus * mask[:, 0:1, :])
+            pool2 = self.pool(x + self._minus * mask[:, 1:2, :])
+            pool3 = self.pool(x + self._minus * mask[:, 2:3, :])
+
+            x = torch.cat([pool1, pool2, pool3], 1)
+            # x = x.squeeze(-1)
+            return  x

train_nyt10_cnn_att.py

@@ -18,7 +18,7 @@
                                              kernel_size=3,
                                              padding_size=1,
                                              word2vec=word2vec,
-                                             dropout=0.5)
+                                             dropout=0.0)
 model = opennre.model.BagAttention(sentence_encoder, len(rel2id), rel2id)
 framework = opennre.framework.BagRE(
     train_path='benchmark/nyt10/nyt10_train.txt',
@@ -32,7 +32,7 @@
     weight_decay=0,
     opt='sgd')
 # Train
-framework.train_model()
+# framework.train_model()
 # Test
 framework.load_state_dict(torch.load(ckpt)['state_dict'])
 result = framework.eval_model(framework.test_loader)

train_nyt10_cnn_att_fixbag.py

@@ -18,7 +18,7 @@
                                              kernel_size=3,
                                              padding_size=1,
                                              word2vec=word2vec,
-                                             dropout=0.5)
+                                             dropout=0.0)
 model = opennre.model.BagAttention(sentence_encoder, len(rel2id), rel2id)
 framework = opennre.framework.BagRE(
     train_path='benchmark/nyt10/nyt10_train.txt',

train_nyt10_pcnn_att.py

@@ -18,16 +18,16 @@
                                              kernel_size=3,
                                              padding_size=1,
                                              word2vec=word2vec,
-                                             dropout=0.5)
-model = nrekit.model.BagAttention(sentence_encoder, len(rel2id), rel2id)
-framework = nrekit.framework.BagRE(
+                                             dropout=0.0)
+model = opennre.model.BagAttention(sentence_encoder, len(rel2id), rel2id)
+framework = opennre.framework.BagRE(
     train_path='benchmark/nyt10/nyt10_train.txt',
     val_path='benchmark/nyt10/nyt10_val.txt',
     test_path='benchmark/nyt10/nyt10_test.txt',
     model=model,
     ckpt=ckpt,
     batch_size=160,
-    max_epoch=60,
+    max_epoch=25,
     lr=0.5,
     weight_decay=0,
     opt='sgd')

train_nyt10_pcnn_att_fixbag.py

@@ -0,0 +1,40 @@
+# coding:utf-8
+import torch
+import numpy as np
+import json
+import opennre
+from opennre import encoder, model, framework
+
+ckpt = 'ckpt/nyt10_pcnn_att_fixbag.pth.tar'
+word2id = json.load(open('pretrain/glove/glove.6B.50d_word2id.json'))
+word2vec = np.load('pretrain/glove/glove.6B.50d_mat.npy')
+rel2id = json.load(open('benchmark/nyt10/nyt10_rel2id.json'))
+sentence_encoder = opennre.encoder.PCNNEncoder(token2id=word2id,
+                                             max_length=120,
+                                             word_size=50,
+                                             position_size=5,
+                                             hidden_size=230,
+                                             blank_padding=True,
+                                             kernel_size=3,
+                                             padding_size=1,
+                                             word2vec=word2vec,
+                                             dropout=0.0)
+model = opennre.model.BagAttention(sentence_encoder, len(rel2id), rel2id)
+framework = opennre.framework.BagRE(
+    train_path='benchmark/nyt10/nyt10_train.txt',
+    val_path='benchmark/nyt10/nyt10_val.txt',
+    test_path='benchmark/nyt10/nyt10_test.txt',
+    model=model,
+    ckpt=ckpt,
+    batch_size=160,
+    max_epoch=60,
+    lr=0.5,
+    weight_decay=0,
+    opt='sgd',
+    bag_size=3)
+# Train
+framework.train_model()
+# Test
+framework.load_state_dict(torch.load(ckpt)['state_dict'])
+result = framework.eval_model(framework.test_loader)
+print('AUC on test set: {}'.format(result['auc']))