Albert pooling dimension mismatches
Author: williamsdaniel888
Pull request: no
Labels: name: wontfix

# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...): **albert-base-v2**

Language I am using the model on (English, Chinese ...): **English**

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Run the following code snippet. You will need the **gdown** package. Note that I am using a modified version of modeling_albert.py to avoid the bug described [here](https://github.com/huggingface/transformers/pull/4095) and [here](https://github.com/huggingface/transformers/issues/1188).

```shell
# Install Transformers
pip install transformers

# Create a working folder (skip if using Colab)
mkdir /content/

# Clone BAT repository to the folder 
git clone https://github.com/akkarimi/Adversarial-Training-for-ABSA.git

# Download patches for BAT scripts in src/
cd /content/
git clone https://github.com/williamsdaniel888/reimagined-octo-tribble.git
mv /content/reimagined-octo-tribble/modeling_albert.py /usr/local/lib/python3.6/dist-packages/transformers/modeling_albert.py
mv /content/reimagined-octo-tribble/absa_data_utils.py /content/Adversarial-Training-for-ABSA/src/absa_data_utils.py
mv /content/reimagined-octo-tribble/asc_bert_pt.py /content/Adversarial-Training-for-ABSA/src/asc_bert_pt.py
mv /content/reimagined-octo-tribble/bat_asc.py /content/Adversarial-Training-for-ABSA/src/bat_asc.py
mv /content/reimagined-octo-tribble/run_asc.py /content/Adversarial-Training-for-ABSA/src/run_asc.py

# Training
cd /content/Adversarial-Training-for-ABSA/script
chmod +x ./run_absa.sh
./run_absa.sh asc laptop_pt laptop pt_asc 1 0
```

2. Open **/content/Adversarial-Training-for-ABSA/run/pt_asc/laptop/1/train_log.txt** to observe the stack trace.

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

**Stack trace:**
```
Traceback (most recent call last):
  File "../src/run_asc.py", line 280, in <module>
    main()
  File "../src/run_asc.py", line 275, in main
    train(args)
  File "../src/run_asc.py", line 122, in train
    _loss, adv_loss = model(input_ids, segment_ids, input_mask, label_ids)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/content/Adversarial-Training-for-ABSA/src/bat_asc.py", line 20, in forward
    _loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/loss.py", line 932, in forward
    ignore_index=self.ignore_index, reduction=self.reduction)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py", line 2317, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py", line 2113, in nll_loss
    .format(input.size(0), target.size(0)))
ValueError: Expected input batch_size (1600) to match target batch_size (16).
```

A diagnostic printout of intermediate variables' sizes within the model's forward() method:
```
Shape of input_ids: torch.Size([16, 100])
Shape of embedding_output: torch.Size([16, 100, 128])
Shape of last encoded layer: torch.Size([16, 100, 768])
Shape of pooled_output: torch.Size([16, 100, 768])
```

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
The pooling unit should pool over the second dimension of the last encoded layer, so that the shape of pooled_output is ```torch.Size([16, 768])``` instead of ```torch.Size([16, 100, 768])```.

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 3.0.0
- Platform: Linux-4.19.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.5.1+cu101 (True)
- Tensorflow version (GPU?): 2.2.0 (True)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

## Possible fix
I have worked around this issue by defining an AlbertPooler class for AlbertModel, similar to the BertPooler method for BertModel in modeling_bert.py:
```python
class AlbertPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output      
```

Make sure to then update **AlbertModel's init()** method:
```
self.pooler = AlbertPooler(config)
```
and **AlbertModel's forward()** method:
```python
pooled_output = self.pooler(sequence_output)
```


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.