Fix workers creation

Author: ivan-chai

hotpp/calibrate.py

@@ -5,16 +5,15 @@
 import torch
 from omegaconf import OmegaConf
 
-from hotpp.data import ShuffledDistributedDataset, DEFAULT_PARALLELIZM
+from hotpp.data import ShuffledDistributedDataset, DEFAULT_PARALLELIZM, get_default_loader_params
 from hotpp.data.module import HotppSampler
 from tqdm import tqdm
 
 logger = logging.getLogger(__name__)
 
 
 def get_loader(dm):
-    loader_params = {"drop_last": False,
-                     "pin_memory": torch.cuda.is_available()}
+    loader_params = get_default_loader_params()
     loader_params.update(dm.train_loader_params)
     dataset = ShuffledDistributedDataset(dm.val_data, rank=None, world_size=None,
                                          num_workers=loader_params.get("num_workers", 0),

hotpp/data/__init__.py

@@ -1,3 +1,3 @@
 from .dataset import HotppDataset, ShuffledDistributedDataset, DEFAULT_PARALLELIZM
-from .module import HotppDataModule
+from .module import HotppDataModule, get_default_loader_params
 from .padded_batch import PaddedBatch

hotpp/data/module.py

@@ -1,17 +1,31 @@
-import torch
+import multiprocessing as mp
 import pytorch_lightning as pl
+import torch
 from .dataset import HotppDataset, ShuffledDistributedDataset, DEFAULT_PARALLELIZM
 
 
 def pop_loader_params(params):
     loader_params = {}
     for key in ["seed", "num_workers", "batch_size", "cache_size", "parallelize", "drop_last", "prefetch_factor",
-                "persistent_workers", "multiprocessing_context"]:
+                "pin_memory", "persistent_workers", "multiprocessing_context"]:
         if key in params:
             loader_params[key] = params.pop(key)
     return loader_params
 
 
+def get_default_loader_params():
+    default_loader_params = {
+        "persistent_workers": True,
+        "pin_memory": torch.cuda.is_available()
+    }
+    available_contexts = mp.get_all_start_methods()
+    for context in ["forkserver", "spawn", "fork"]:
+        if context in available_contexts:
+            default_loader_params["multiprocessing_context"] = context
+            break
+    return default_loader_params
+
+
 class HotppSampler(torch.utils.data.DistributedSampler):
     def __init__(self, dataset):
         # Skip super init.
@@ -106,10 +120,8 @@ def splits(self):
     def train_dataloader(self, rank=None, world_size=None):
         rank = self.trainer.global_rank if rank is None else rank
         world_size = self.trainer.world_size if world_size is None else world_size
-        loader_params = {"drop_last": True,
-                         "multiprocessing_context": "spawn",
-                         "persistent_workers": True,
-                         "pin_memory": torch.cuda.is_available()}
+        loader_params = get_default_loader_params()
+        loader_params.update({"drop_last": True})
         loader_params.update(self.train_loader_params)
         dataset = ShuffledDistributedDataset(self.train_data, rank=rank, world_size=world_size,
                                              cache_size=loader_params.pop("cache_size", 4096),
@@ -127,9 +139,7 @@ def train_dataloader(self, rank=None, world_size=None):
     def val_dataloader(self, rank=None, world_size=None):
         rank = self.trainer.global_rank if rank is None else rank
         world_size = self.trainer.world_size if world_size is None else world_size
-        loader_params = {"multiprocessing_context": "spawn",
-                         "persistent_workers": True,
-                         "pin_memory": torch.cuda.is_available()}
+        loader_params = get_default_loader_params()
         loader_params.update(self.val_loader_params)
         dataset = ShuffledDistributedDataset(self.val_data, rank=rank, world_size=world_size,
                                              parallelize=loader_params.pop("parallelize", DEFAULT_PARALLELIZM))  # Disable shuffle.
@@ -143,9 +153,7 @@ def val_dataloader(self, rank=None, world_size=None):
     def test_dataloader(self, rank=None, world_size=None):
         rank = self.trainer.global_rank if rank is None else rank
         world_size = self.trainer.world_size if world_size is None else world_size
-        loader_params = {"multiprocessing_context": "spawn",
-                         "persistent_workers": True,
-                         "pin_memory": torch.cuda.is_available()}
+        loader_params = get_default_loader_params()
         loader_params.update(self.test_loader_params)
         dataset = ShuffledDistributedDataset(self.test_data, rank=rank, world_size=world_size,
                                              parallelize=loader_params.pop("parallelize", DEFAULT_PARALLELIZM))  # Disable shuffle.

hotpp/embed.py

@@ -13,7 +13,7 @@
 from torchmetrics.utilities.distributed import gather_all_tensors
 
 from .common import get_trainer
-from .data import ShuffledDistributedDataset, DEFAULT_PARALLELIZM
+from .data import ShuffledDistributedDataset, DEFAULT_PARALLELIZM, get_default_loader_params
 
 logger = logging.getLogger(__name__)
 
@@ -125,7 +125,8 @@ def test_dataloader(self):
             collate_fn=dataset.dataset.collate_fn,
             shuffle=False,
             num_workers=num_workers,
-            batch_size=loader_params.get("batch_size", 1)
+            batch_size=loader_params.get("batch_size", 1),
+            **get_default_loader_params()
         )