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):
