BertForTokenClassification, logits on sequence tagging task 
Author: mbkrafft
Pull request: no
Labels: name: wontfix

# ❓ Questions & Help

## Details
<!-- Description of your issue -->
I'm trying to use BertForTokenClassification to do a sequence tagging task. During training, everything seems fine and loss is decreasing per epoch, but when I put the model in evaluation mode and feed it inputs, the logits for each token are all the same. 

My code looks like this:
```
class PretrainedBert(nn.Module):

    def __init__(self,
                 device,
                 pretrained_path='bert-base-multilingual-cased',
                 output_dim=5,
                 learning_rate=0.01,
                 eps=1e-8,
                 max_grad_norm=1.0,
                 model_name='bert_multilingual',
                 cased=True
                 ):
        super(PretrainedBert, self).__init__()

        self.device = device
        self.output_dim = output_dim
        self.learning_rate = learning_rate
        self.max_grad_norm = max_grad_norm
        self.model_name = model_name
        self.cased = cased

        # Load pretrained bert
        self.bert = BertForTokenClassification.from_pretrained(pretrained_path, num_labels=output_dim, output_attentions = False, output_hidden_states = False)
        self.tokenizer = BertTokenizer.from_pretrained(pretrained_path, do_lower_case=not cased)
        
        param_optimizer = list(self.bert.named_parameters())
        no_decay = ['bias', 'gamma', 'beta']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
             'weight_decay_rate': 0.01},
            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
             'weight_decay_rate': 0.0}]
        self.optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate, eps=eps)
        
        
    def forward(self, input_ids, attention_mask, labels=None):
        output = self.bert(input_ids, attention_mask=attention_mask, labels=labels)        
        return output

    
    def fit(self, train_loader, dev_loader, epochs=10):
        """
        :param test_loader: torch.utils.data.DataLoader object
        :param test_loader: torch.utils.data.DataLoader object
        :param epochs: int
        """
        scheduler = get_linear_schedule_with_warmup(self.optimizer, num_warmup_steps=0,num_training_steps=len(train_loader)*epochs)

        for epoch in range(epochs):
            # Iterate over training data
            epoch_time = time.time()
            self.train()
            epoch_loss = 0
            num_batches = 0

            for raw_text, x, y, idx in tqdm(train_loader):                
                self.zero_grad()

                # Get gold labels and extend them to cover the wordpieces created by bert
                y = self._extend_labels(raw_text, y)
                y, _ = pad_packed_sequence(y, batch_first=True)
                batches_len, seq_length = y.size()
                y.to(self.device)

                # Run input through bert and get logits
                bert_tokenized = self.tokenizer.batch_encode_plus([' '.join(
                    text) for text in raw_text], max_length=seq_length, pad_to_max_length=True)
                
                input_ids, token_type_ids, attention_masks = bert_tokenized[
                    'input_ids'], bert_tokenized['token_type_ids'], bert_tokenized['attention_mask']
                
                loss, logits = self.forward(torch.LongTensor(
                    input_ids), torch.LongTensor(attention_masks), labels=y)    
                
                # Help prevent the "exploding gradients" problem.
                torch.nn.utils.clip_grad_norm_(parameters=self.bert.parameters(), max_norm=self.max_grad_norm)
                
                # update parameters and learning rate
                loss.backward()
                self.optimizer.step()
                scheduler.step()
                
                epoch_loss += loss.item()
                num_batches += 1

            epoch_time = time.time() - epoch_time

            print()
            print("Epoch {0} loss: {1:.3f}".format(epoch + 1,
                                                   epoch_loss / num_batches))

            print("Dev")
            binary_f1, propor_f1 = self.evaluate(dev_loader)


    def predict(self, test_loader):
        """
        :param test_loader: torch.utils.data.DataLoader object with
                            batch_size=1
        """
        self.eval()
        predictions, golds, sents = [], [], []

        for raw_text, x, y, idx in tqdm(test_loader):     
            # Run input through bert and get logits
            bert_tokenized = self.tokenizer.encode_plus(' '.join(raw_text[0]))

            input_ids, token_type_ids, attention_masks = bert_tokenized[
                'input_ids'], bert_tokenized['token_type_ids'], bert_tokenized['attention_mask']

            with torch.no_grad():
                logits = self.forward(torch.LongTensor(
                    input_ids).unsqueeze(0), torch.LongTensor(attention_masks).unsqueeze(0)) 
   
            # remove batch dim and [CLS]+[SEP] tokens from logits
            logits = logits[0].squeeze(0)[1:-1:]
	   
            # mean pool wordpiece rows of logits and argmax to get predictions 
            preds = self._logits_to_preds(logits, y[0], raw_text[0])

            predictions.append(preds)
            golds.append(y[0])
            sents.append(raw_text[0])

        return predictions, golds, sents

    
    def evaluate(self, test_loader):
        """
        Returns the binary and proportional F1 scores of the model on the examples passed via test_loader.
        :param test_loader: torch.utils.data.DataLoader object with
                            batch_size=1
        """
        preds, golds, sents = self.predict(test_loader)

        
        flat_preds = [int(i) for l in preds for i in l]
        flat_golds = [int(i) for l in golds for i in l]

        analysis = get_analysis(sents, preds, golds)
        binary_f1 = binary_analysis(analysis)
        propor_f1 = proportional_analysis(flat_golds, flat_preds)
        
        return binary_f1, propor_f1


    def _extend_labels(self, text, labels):
        extended_labels = []

        for idx, (text_seq, label_seq) in enumerate(zip(text, labels)):
            extended_seq_labels = []

            for word, label in zip(text_seq, label_seq):
                n_subwords = len(self.tokenizer.tokenize(word))
                extended_seq_labels.extend([label.item()] * n_subwords)

            extended_labels.append(torch.LongTensor(extended_seq_labels))

        extended_labels = pack_sequence(extended_labels, enforce_sorted=False)
        
        return extended_labels

    
    def _logits_to_preds(self, logits, y, raw_text):
        preds = torch.zeros(len(y))
        head = 0
        for idx, word in enumerate(raw_text):
            n_subwords = len(self.tokenizer.tokenize(word))

            if n_subwords > 1:
                preds[idx] = torch.argmax(torch.mean(
                    logits[head:head+n_subwords, :], dim=0))
            else:
                preds[idx] = torch.argmax(logits[head, :])

            head = head + n_subwords

        return preds
```

Example input:

```
raw_text = ['Donkey', 'Kong', 'Country', ':']
y =  [tensor([0, 0, 0, 0])]
```

Produces these results:
```
tokenized_words = ['Don', '##key', 'Kong', 'Country', ':']

logits = tensor([[-3.2811,  1.1715,  1.2381,  0.5201,  1.0921],
        [-3.2813,  1.1715,  1.2382,  0.5201,  1.0922],
        [-3.2815,  1.1716,  1.2383,  0.5202,  1.0923],
        [-3.2814,  1.1716,  1.2383,  0.5202,  1.0922],
        [-3.2811,  1.1715,  1.2381,  0.5201,  1.0921]])

preds = tensor([2., 2., 2., 2.])
```


Any ideas for what I'm doing wrong? Thanks for any help in advance. 


Comments:
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.