Feature extraction pipeline increasing memory use
Author: quancore
Pull request: no

### System Info

UBUNTU 22.04

### Who can help?

@Narsil 

### Reproduction

```python
import os
import argparse

import hickle as hkl
import numpy as np
import pandas as pd
from sklearn import model_selection
from transformers import pipeline, AutoTokenizer, AutoModel
import torch
from torch.utils.data import Dataset

from tqdm import tqdm


def is_file_path(path):
    if os.path.isfile(path):
        return path
    else:
        raise argparse.ArgumentTypeError(f"{path} is not a valid file path")


def is_dir_path(path):
    if os.path.isdir(path):
        return path
    else:
        raise argparse.ArgumentTypeError(f"{path} is not a valid dir path")


parser = argparse.ArgumentParser()

parser.add_argument("-d", "--dataframe", help="Path of the dataframe", type=is_file_path, dest="df_path")
parser.add_argument("-ed", "--embedding_dir", help="Path of embeddings stored", type=is_dir_path, dest="embedding_dir")
parser.add_argument("-V", "--version", help="show program version", action="store_true", dest="version")
args = parser.parse_args()

MODEL_NAME = "anferico/bert-for-patents"

pipe = None
test_size = 1000
input_text_column = "text"
label_column = "SUBFIELD_ID"


class ListDataset(Dataset):
    def __init__(self, original_list):
        self.original_list = original_list
    def __len__(self):
        return len(self.original_list)
    def __getitem__(self, i):
        return self.original_list[i]


def get_pipeline(device_index):
    global pipe
    if pipe is None:
        print(f"Creating pipeline for device: {device_index}")
        tokenizer = AutoTokenizer.from_pretrained(
            MODEL_NAME, do_lower_case=True, model_max_length=512, truncation=True, padding=True, pad_to_max_length=True
        )
        device = torch.device(f"cuda:{device_index}" if torch.cuda.is_available() else "cpu")
        model = AutoModel.from_pretrained(MODEL_NAME).to(device)
        local_pipe = pipeline(
            "feature-extraction",
            model=model,
            tokenizer=tokenizer,
            max_length=512,
            truncation=True,
            padding=True,
            pad_to_max_length=True,
            device=device,
            framework="pt",
            batch_size=16,
        )
        pipe = local_pipe
    return pipe


def extract_embeddings(df, col, mode="cls"):
    dataset = ListDataset(df[col].tolist())
    if mode == "cls":
        pipe = get_pipeline(0)
        embeddings = []
        for embedding in tqdm(pipe(dataset, max_length=512, truncation=True, num_workers=8), total=len(dataset)):
            embeddings.append(np.squeeze(embedding)[0])

        return embeddings


def main():
    df = pd.read_csv(os.path.join(args.df_path), sep='\t', header=0)
    train_df, test_df = model_selection.train_test_split(df, shuffle=True, stratify=df[label_column],
                                                         train_size=len(df) - test_size, random_state=50)
    print("Extracting train embeddings...")
    x_train_bert = extract_embeddings(train_df, input_text_column)
    print("Extracting test embeddings...")
    x_test_bert = extract_embeddings(test_df, input_text_column)


if __name__ == "__main__":
    main()
```


### Expected behavior

I am trying to extract text embeddings from a text column of around 65k entries. I need to store the embeddings in the memory t train the downstream sklearn classifier (no online, incremental mode). The memory use of the above feature extraction gets insane like 128 GB, may I ask why this is the case (I am just storing one NumPy array of 1024 float numbers per entry, so it has to be much much smaller).

Comments:
Hi @quancore ,

Is this line normal ? 
```
    dataset = ListDataset(df[col].tolist())
```
This will load the entire dataset in memory, not sure if 65k entries are big or not but they can add up fast if they are sizeable documents.
`    df = pd.read_csv(os.path.join(args.df_path), sep='\t', header=0)` 

As far as I understand you're loading the entire file here (so larger than just the col entries you're looking for)


Finally you're setting `max_length=512`  which means your larges embedding is `512 * 1024` x `65 000` that's roughly `30Go` .
In `transformers` there's no reduction of the embedding which is done sometimes by `sentence-transformers` (either looking at the embedding of the first token, averaging, maxing or other reduction mecanisms)
Could this part be missing ?

`pad_to_max_length=True,` means all the tensors will be 512 x 1024.

Also I'm not sure, but I think the output of the pipeline by default is a raw list meaning it will take up more space than it's `numpy` equivalent. You could try using the the new `return_tensors=True` parameter to receive directly the embedding in tensor format.

Tell me if any  of this helps solve your use case !


HI @Narsil , thank you for the answer.
```
Is this line normal ?

dataset = ListDataset(df[col].tolist())

This will load the entire dataset in memory, not sure if 65k entries are big or not but they can add up fast if they are sizeable documents.
df = pd.read_csv(os.path.join(args.df_path), sep='\t', header=0)
```
The dataset is not that big and ListDataset instance example actually comes from one of the answers on the forum: https://discuss.huggingface.co/t/progress-bar-for-hf-pipelines/20498/2
That's why I am loading the data frame and converting it to a list for use, do you have a better suggestion while seeing the progress?

```
Finally you're setting max_length=512 which means your larges embedding is 512 * 1024 x 65 000 that's roughly 30Go .
In transformers there's no reduction of the embedding which is done sometimes by sentence-transformers (either looking at the embedding of the first token, averaging, maxing or other reduction mecanisms)
Could this part be missing?
```
As you have calculated, it should be around 30GB  or around, but it is reaching up to 128 GB which is insane. Right now, I am only interested in the CLS token, so if I set the max token to 1, will it give me only the CLS token?

```
Also I'm not sure, but I think the output of the pipeline by default is a raw list meaning it will take up more space than it's numpy equivalent. You could try using the the new return_tensors=True parameter to receive directly the embedding in tensor format.
```
I will try this.

> The dataset is not that big and ListDataset instance example actually comes from one of the answers on the forum: https://discuss.huggingface.co/t/progress-bar-for-hf-pipelines/20498/2

If the dataset is not that big it will work fine, but since you already have the data you don't need to do `.tolist()`, I would do something like

```python
class MyDataset:
   def __init__(self, panda_frame, col):
       self.panda_fram = panda_frame
       self.col = col
   def __len__(self):
        return len(self.panda_frame[self.col])
   def __getitem__(self, i):
        return self.panda_frame[self.col][i]
 ```
I haven't tested it, but with this gist you could get away without copying anything.


> As you have calculated, it should be around 30GB or around, but it is reaching up to 128 GB which is insane. Right now, I am only interested in the CLS token, so if I set the max token to 1, will it give me only the CLS token?

Don't  use `pad_to_max_length` imo. This is unnecessary in a lot of cases. What you want to do is to run the model on the actual full sentence, but keep around only the embedding for the first token (which should be token[0], but I don't know the model you are using, it may not exist depending on how it's setup.)


```
            embeddings.append(embedding[0]) # Embedding should be `seq_len, hidden_dim` so `embedding[0] should be `hidden_dim`.
 ```
 
 Also since you are in a pipeline, you could also write to disk the results, in a dataset, a different file for each embeddings or something like that. Then you could run the entire pipeline on very little memory, that's basically the whole point of pipeline, to try and limit aggressively the memory necessary. 
 
For 30Go and 128Go, it's hard to answer exactly without checking, but constants do pile up fast in Python. Anything reference another viable will keep the whole data alive for instance.
The pipeline itself shouldn't keep anything around but `lists` are much more wasteful than tensors memory wise, also at those levels of memory there's probably a big chunk of fragmentation of memory which might increase the overall usage. Garbage collection might not be able to keep up with the amount of data you generate etc, etc..

If you could reproduce an issue on a smaller scale example using something  like `busybox time ...` to showcase that we're using too much memory for a given task (can't really reproduce your example right now) that'd be lovely.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.