Skip to content

Instantly share code, notes, and snippets.

@leriomaggio
Last active April 11, 2021 11:34
Show Gist options
  • Save leriomaggio/ee7444d8cfd36e7ed9bb6c596089ac4b to your computer and use it in GitHub Desktop.
Save leriomaggio/ee7444d8cfd36e7ed9bb6c596089ac4b to your computer and use it in GitHub Desktop.
Python and Random State Sharing with NumPy and PyTorch, depending on different multiprocessing start method (i.e. fork vs spawn)
import numpy as np
import multiprocessing as mp
from argparse import ArgumentParser
def get_current_seed(wid):
s = np.random.get_state()[1][0]
print(f"Pool Worker {wid}: Random State {s}")
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument(
"--with-fork",
action="store_true",
help="Force fork method",
)
args = parser.parse_args()
if args.with_fork:
mp.set_start_method("fork")
print("Start Method: ", mp.get_start_method())
SEED = 123456
# set seed
np.random.seed(SEED)
print("Main Process : Random State", np.random.get_state()[1][0])
with mp.Pool(5) as p:
p.map(get_current_seed, range(5))
# Identical to NumPy version, but using PyTorch random
import torch as th
import multiprocessing as mp
from argparse import ArgumentParser
def get_current_seed(wid):
s = th.random.initial_seed()
print(f"Pool Worker {wid}: Random State {s}")
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument(
"--with-fork",
action="store_true",
help="Force fork method",
)
args = parser.parse_args()
if args.with_fork:
mp.set_start_method("fork")
print("Start Method: ", mp.get_start_method())
SEED = 123456
# set seed
th.random.manual_seed(SEED)
print("Main Process : Random State", th.random.initial_seed())
with mp.Pool(5) as p:
p.map(get_current_seed, range(5))

Replicability

This output has been generated using Python 3.8 on macOS, which is the first release of Py3 on macOS to use spawn as the default start_method (more here)

numpy

❯ python multiprocessing_seed_numpy.py
Start Method:  spawn
Main Process : Random State 123456
Pool Worker 0: Random State 2147483648
Pool Worker 1: Random State 2147483648
Pool Worker 3: Random State 2147483648
Pool Worker 2: Random State 2147483648
Pool Worker 4: Random State 2147483648

❯ python multiprocessing_seed_numpy.py --with-fork
Start Method:  fork
Main Process : Random State 123456
Pool Worker 0: Random State 123456
Pool Worker 1: Random State 123456
Pool Worker 2: Random State 123456
Pool Worker 3: Random State 123456
Pool Worker 4: Random State 123456

torch

❯ python multiprocessing_seed_torch.py
Start Method:  spawn
Main Process : Random State 123456
Pool Worker 0: Random State 12008651253415256287
Pool Worker 1: Random State 12008651253415256287
Pool Worker 2: Random State 12008651253415256287
Pool Worker 3: Random State 12008651253415256287
Pool Worker 4: Random State 12008651253415256287

❯ python multiprocessing_seed_torch.py --with-fork
Start Method:  fork
Main Process : Random State 123456
Pool Worker 0: Random State 123456
Pool Worker 1: Random State 123456
Pool Worker 2: Random State 123456
Pool Worker 3: Random State 123456
Pool Worker 4: Random State 123456
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment