def _worker_loop_seed(worker_id, dataset, index_queue, data_queue, collate_fn, seed, worker_init_fn, worker_init_args, worker_init_kwargs): reset_global_seed(seed) if worker_init_fn is not None: worker_init_fn(worker_id, *worker_init_args, **worker_init_kwargs) _worker_loop(dataset, index_queue, data_queue, collate_fn)
def __call__(self, worker_id): seed = (self._base_seed + worker_id) % 42964967296 reset_global_seed(seed) if self._fn_init is not None: args = self._args[worker_id] kwargs = self._kwargs[worker_id] self._fn_init(worker_id, *args, **kwargs)
def __call__(self, worker_id): seed = (self._base_seed + worker_id) % 42964967296 reset_global_seed(seed) if self._fn_init is not None: args = self._args[worker_id] kwargs = self._kwargs[worker_id] self._fn_init(worker_id, *args, **kwargs) if self._fn_recv is not None: if self._pipe_master is not None and len( self._pipe_master.queues) > 0: self._fn_recv.worker_init(self._pipe_master.queues[worker_id])
args = parser.parse_args() args.use_gpu = args.use_gpu and torch.cuda.is_available() if args.dump_dir is not None: io.mkdir(args.dump_dir) args.log_file = os.path.join(args.dump_dir, 'log.log') set_output_file(args.log_file) else: args.checkpoints_dir = None args.summary_file = None if args.seed is not None: import jacinle.random as random random.reset_global_seed(args.seed) args.task_is_outdegree = args.task in ['outdegree'] args.task_is_connectivity = args.task in ['connectivity'] args.task_is_adjacent = args.task in ['adjacent', 'adjacent-mnist'] args.task_is_family_tree = args.task in [ 'has-father', 'has-sister', 'grandparents', 'uncle', 'maternal-great-uncle' ] args.task_is_mnist_input = args.task in ['adjacent-mnist'] args.task_is_1d_output = args.task in [ 'outdegree', 'adjacent', 'adjacent-mnist', 'has-father', 'has-sister' ] class LeNet(nn.Module): def __init__(self):