mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
dcd1216efe
Summary: Fixes https://github.com/pytorch/pytorch/issues/28389 Intel's OpenMP implementation sets the thread affinity on the first call to an OpenMP function after a fork. By adding an atfork handler we can force this to happen before a user tries to set the affinity in their own DataLoader `worker_init_fn`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/29006 Differential Revision: D18782456 Pulled By: ezyang fbshipit-source-id: ce0b515256da0cf18ceb125e0cdec99a3311bbd3
32 lines
805 B
Python
32 lines
805 B
Python
from __future__ import absolute_import
|
|
import sys
|
|
|
|
__all__ = ['register_after_fork']
|
|
|
|
if sys.version_info < (3, 7):
|
|
import multiprocessing.util as _util
|
|
|
|
def _register(func):
|
|
def wrapper(arg):
|
|
func()
|
|
_util.register_after_fork(_register, wrapper)
|
|
else:
|
|
import os
|
|
|
|
def _register(func):
|
|
os.register_at_fork(after_in_child=func)
|
|
|
|
def register_after_fork(func):
|
|
"""Register a callable to be executed in the child process after a fork.
|
|
|
|
Note:
|
|
In python < 3.7 this will only work with processes created using the
|
|
``multiprocessing`` module. In python >= 3.7 it also works with
|
|
``os.fork()``.
|
|
|
|
Arguments:
|
|
func (function): Function taking no arguments to be called in the child after fork
|
|
|
|
"""
|
|
_register(func)
|