Revert "fix mp in DataLoader (#2506) (#2507)" (#2521)

This reverts commit 5576e6f.
wenet-e2e · May 8, 2024 · 2258c72 · 2258c72
1 parent bd22fae
commit 2258c72
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 27 deletions.
diff --git a/test/wenet/dataset/test_datapipes.py b/test/wenet/dataset/test_datapipes.py
@@ -2,7 +2,6 @@
 import torch
 from torch.utils.data import datapipes
 from torch.utils.data.datapipes.iter import IterableWrapper
-import torch.multiprocessing as mp
 from functools import partial
 
 from wenet.dataset.datapipes import (RepeatDatapipe, SortDataPipe,
@@ -109,11 +108,9 @@ def test_dynamic_batch_datapipe(data_list):
         window_class=DynamicBatchWindow(max_frames_in_batch),
         wrapper_class=padding)
 
-    dataloader = torch.utils.data.DataLoader(
-        dataset,
-        batch_size=None,
-        num_workers=2,
-        multiprocessing_context=mp.get_context("spawn"))
+    dataloader = torch.utils.data.DataLoader(dataset,
+                                             batch_size=None,
+                                             num_workers=2)
     for d in dataloader:
         assert d['feats'].size(1) <= max_frames_in_batch
 

diff --git a/test/wenet/dataset/test_dataset.py b/test/wenet/dataset/test_dataset.py
@@ -1,6 +1,5 @@
 import pytest
 import torch
-import torch.multiprocessing as mp
 from wenet.dataset.dataset import Dataset
 from wenet.text.char_tokenizer import CharTokenizer
 
@@ -55,11 +54,9 @@ def test_dataset(params):
                       data_list,
                       tokenizer=tokenizer,
                       conf=dataset_conf)
-    dataloader = torch.utils.data.DataLoader(
-        dataset,
-        batch_size=None,
-        num_workers=4,
-        persistent_workers=True,
-        multiprocessing_context=mp.get_context("spawn"))
+    dataloader = torch.utils.data.DataLoader(dataset,
+                                             batch_size=None,
+                                             num_workers=4,
+                                             persistent_workers=True)
     for d in dataloader:
         pass
diff --git a/tools/compute_cmvn_stats.py b/tools/compute_cmvn_stats.py
@@ -11,7 +11,6 @@
 import torchaudio
 import torchaudio.compliance.kaldi as kaldi
 from torch.utils.data import Dataset, DataLoader
-import torch.multiprocessing as mp
 
 
 class CollateFunc(object):
@@ -108,14 +107,12 @@ def __getitem__(self, idx):
     collate_func = CollateFunc(feat_dim, resample_rate)
     dataset = AudioDataset(args.in_scp)
     batch_size = 20
-    mp_context = mp.get_context("spawn") if args.num_workers > 0 else None
     data_loader = DataLoader(dataset,
                              batch_size=batch_size,
                              shuffle=True,
                              sampler=None,
                              num_workers=args.num_workers,
-                             collate_fn=collate_func,
-                             multiprocessing_context=mp_context)
+                             collate_fn=collate_func)
 
     with torch.no_grad():
         all_number = 0

diff --git a/wenet/bin/recognize.py b/wenet/bin/recognize.py
@@ -22,7 +22,6 @@
 import torch
 import yaml
 from torch.utils.data import DataLoader
-import torch.multiprocessing as mp
 
 from wenet.dataset.dataset import Dataset
 from wenet.utils.config import override_config
@@ -223,11 +222,9 @@ def main():
                            test_conf,
                            partition=False)
 
-    mp_context = mp.get_context("spawn") if args.num_workers > 0 else None
     test_data_loader = DataLoader(test_dataset,
                                   batch_size=None,
-                                  num_workers=args.num_workers,
-                                  multiprocessing_context=mp_context)
+                                  num_workers=args.num_workers)
 
     # Init asr model from configs
     args.jit = False

diff --git a/wenet/utils/train_utils.py b/wenet/utils/train_utils.py
@@ -26,7 +26,6 @@
 
 import torch.optim as optim
 import torch.distributed as dist
-import torch.multiprocessing as mp
 
 from tensorboardX import SummaryWriter
 from torch.utils.data import DataLoader
@@ -346,23 +345,20 @@ def init_dataset_and_dataloader(args, configs, tokenizer, seed=777):
 
     # NOTE(xcsong): Why we prefer persistent_workers=True ?
     #   https://discuss.pytorch.org/t/what-are-the-dis-advantages-of-persistent-workers/102110
-    mp_context = mp.get_context("spawn") if args.num_workers > 0 else None
     train_data_loader = DataLoader(train_dataset,
                                    batch_size=None,
                                    pin_memory=args.pin_memory,
                                    num_workers=args.num_workers,
                                    persistent_workers=True,
                                    generator=generator,
-                                   prefetch_factor=args.prefetch,
-                                   multiprocessing_context=mp_context)
+                                   prefetch_factor=args.prefetch)
     cv_data_loader = DataLoader(cv_dataset,
                                 batch_size=None,
                                 pin_memory=args.pin_memory,
                                 num_workers=args.num_workers,
                                 persistent_workers=True,
                                 generator=generator,
-                                prefetch_factor=args.prefetch,
-                                multiprocessing_context=mp_context)
+                                prefetch_factor=args.prefetch)
     return train_dataset, cv_dataset, train_data_loader, cv_data_loader