Multiprocessing in a `filter` or `map` function with a Pytorch model
Opened by HugoLaurencon (closed)

### Describe the bug

I am trying to use a Pytorch model loaded on CPUs with multiple processes with a `.map` or a `.filter` method.

Usually, when dealing with models that are non-pickable, creating a class such that the `map` function is the method `__call__`, and adding `reduce` helps to solve the problem.

However, here, the command hangs without throwing an error.


### Steps to reproduce the bug

```
from datasets import Dataset
import torch
from torch import nn
from torchvision import models
​
​
class FilterFunction:
    #__slots__ = ("path_model", "model")  # Doesn't change anything uncommented
    def __init__(self, path_model):
        self.path_model = path_model
        model = models.resnet50()
        model.fc = nn.Sequential(
            nn.Linear(2048, 512),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(512, 10),
            nn.LogSoftmax(dim=1)
        )
        model.load_state_dict(torch.load(path_model, map_location=torch.device("cpu")))
        model.eval()
        self.model = model
    def __call__(self, batch):
        return [True] * len(batch["id"])
    # Comment this to have an error
    def __reduce__(self):
        return (self.__class__, (self.path_model,))
​
​
dataset = Dataset.from_dict({"id": [0, 1, 2, 4]})
​
# Download (100 MB) at https://github.com/emiliantolo/pytorch_nsfw_model/raw/master/ResNet50_nsfw_model.pth
path_model = "/fsx/hugo/nsfw_image/ResNet50_nsfw_model.pth"
​
filter_function = FilterFunction(path_model=path_model)
​
# Works
filtered_dataset = dataset.filter(filter_function, num_proc=1, batched=True, batch_size=2)
# Doesn't work
filtered_dataset = dataset.filter(filter_function, num_proc=2, batched=True, batch_size=2)
```

### Expected behavior

The command `filtered_dataset = dataset.filter(filter_function, num_proc=2, batched=True, batch_size=2)` should work and not hang.

### Environment info

Datasets: 2.11.0
Pyarrow: 11.0.0
Ubuntu

Comments:
Hi ! PyTorch may hang when calling `load_state_dict()` in a subprocess. To fix that, set the multiprocessing start method to "spawn". Since `datasets` uses `multiprocess`, you should do:

```python
# Required to avoid issues with pytorch (otherwise hangs during load_state_dict in multiprocessing)
import multiprocess.context as ctx
ctx._force_start_method('spawn')
```

Also make sure to run your main code in `if __name__ == "__main__":` to avoid issues with python multiprocesing

Thanks!

@lhoestq Hello, I also encountered this problem but maybe with another reason. Here is my code:
```python
tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir, model_max_length=training_args.model_max_length)
data = load_dataset("json", data_files=data_args.train_file, cache_dir=data_args.data_cache_dir)
def func(samples):
    # main operation
    for sentence_value in samples:
        sentence_ids = tokenizer.encode(sentence_value, add_special_tokens=False, max_length=tokenizer.model_max_length, truncation=True)
    ... ...
train_data = data["train"].shuffle().map(func, num_proc=os.cpu_count())
```
It hangs after the progress reaches 100%. Could you help me point out the reason?

@SkyAndCloud your issue doesn't seem related to the original post - could you open a new issue and provide more details ? (size of the dataset, number of cpus, how much time it took to run, `datasets` version)

@lhoestq Hi, I just solved this problem. Because the input is extremely long and the tokenizer requests a large amount of memory, which leads to a OOM error and may eventually causes the hang. I didn't filter those too-long sentences because I thought `tokenizer` would stop once the length exceeds the `max_length`. However, it actually firstly complete the tokenization of entire sentence and then truncate it.