Multiprocessing in a `filter` or `map` function with a Pytorch model
MEMBER
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
2023-04-24T10:41:51Z
https://github.com/huggingface/datasets/issues/5786#issuecomment-1519884686
https://api.github.com/repos/huggingface/datasets/issues/5786
IC_kwDODunzps5al5mO
https://api.github.com/repos/huggingface/datasets/issues/comments/1519884686/reactions
2023-04-24T10:41:51Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1519884686
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
Thanks!
2023-04-24T10:43:58Z
https://github.com/huggingface/datasets/issues/5786#issuecomment-1519887372
https://api.github.com/repos/huggingface/datasets/issues/5786
IC_kwDODunzps5al6QM
https://api.github.com/repos/huggingface/datasets/issues/comments/1519887372/reactions
2023-04-24T10:43:58Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1519887372
https://avatars.githubusercontent.com/u/44556846?v=4
https://api.github.com/users/HugoLaurencon/events{/privacy}
https://api.github.com/users/HugoLaurencon/followers
https://api.github.com/users/HugoLaurencon/following{/other_user}
https://api.github.com/users/HugoLaurencon/gists{/gist_id}
https://github.com/HugoLaurencon
HugoLaurencon
MDQ6VXNlcjQ0NTU2ODQ2
https://api.github.com/users/HugoLaurencon/orgs
https://api.github.com/users/HugoLaurencon/received_events
https://api.github.com/users/HugoLaurencon/repos
https://api.github.com/users/HugoLaurencon/starred{/owner}{/repo}
https://api.github.com/users/HugoLaurencon/subscriptions
User
https://api.github.com/users/HugoLaurencon
public
NONE
@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?
2023-05-29T08:37:37Z
https://github.com/huggingface/datasets/issues/5786#issuecomment-1566764875
https://api.github.com/repos/huggingface/datasets/issues/5786
IC_kwDODunzps5dYu9L
https://api.github.com/repos/huggingface/datasets/issues/comments/1566764875/reactions
2023-05-29T08:37:37Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1566764875
https://avatars.githubusercontent.com/u/12583017?v=4
https://api.github.com/users/SkyAndCloud/events{/privacy}
https://api.github.com/users/SkyAndCloud/followers
https://api.github.com/users/SkyAndCloud/following{/other_user}
https://api.github.com/users/SkyAndCloud/gists{/gist_id}
https://github.com/SkyAndCloud
SkyAndCloud
MDQ6VXNlcjEyNTgzMDE3
https://api.github.com/users/SkyAndCloud/orgs
https://api.github.com/users/SkyAndCloud/received_events
https://api.github.com/users/SkyAndCloud/repos
https://api.github.com/users/SkyAndCloud/starred{/owner}{/repo}
https://api.github.com/users/SkyAndCloud/subscriptions
User
https://api.github.com/users/SkyAndCloud
public
MEMBER
@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)
2023-05-30T09:19:01Z
https://github.com/huggingface/datasets/issues/5786#issuecomment-1568085199
https://api.github.com/repos/huggingface/datasets/issues/5786
IC_kwDODunzps5ddxTP
https://api.github.com/repos/huggingface/datasets/issues/comments/1568085199/reactions
2023-05-30T09:19:01Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1568085199
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
@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.
2023-05-30T09:56:30Z
https://github.com/huggingface/datasets/issues/5786#issuecomment-1568144246
https://api.github.com/repos/huggingface/datasets/issues/5786
IC_kwDODunzps5dd_t2
https://api.github.com/repos/huggingface/datasets/issues/comments/1568144246/reactions
2023-05-30T09:56:30Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1568144246
https://avatars.githubusercontent.com/u/12583017?v=4
https://api.github.com/users/SkyAndCloud/events{/privacy}
https://api.github.com/users/SkyAndCloud/followers
https://api.github.com/users/SkyAndCloud/following{/other_user}
https://api.github.com/users/SkyAndCloud/gists{/gist_id}
https://github.com/SkyAndCloud
SkyAndCloud
MDQ6VXNlcjEyNTgzMDE3
https://api.github.com/users/SkyAndCloud/orgs
https://api.github.com/users/SkyAndCloud/received_events
https://api.github.com/users/SkyAndCloud/repos
https://api.github.com/users/SkyAndCloud/starred{/owner}{/repo}
https://api.github.com/users/SkyAndCloud/subscriptions
User
https://api.github.com/users/SkyAndCloud
public
### 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