Function _worker_loop

imperative/python/megengine/data/dataloader.py:838–932

(
    dataset,
    index_queue,
    data_queue,
    done_event,
    transform,
    collate,
    batch_size,
    seed,
    worker_id,
    num_workers,
    datakind,
    parallel_stream,
)

def _worker_loop(
    dataset,
    index_queue,
    data_queue,
    done_event,
    transform,
    collate,
    batch_size,
    seed,
    worker_id,
    num_workers,
    datakind,
    parallel_stream,
):
    _set_worker_signal_handlers()
    random.seed(seed)
    np.random.seed(seed)
    watchdog = ManagerWatchdog()
    iteration_end = False
    fetcher = map_fetcher
    if datakind == "stream":
        global _worker_info
        _worker_info = WorkerInfo(idx=worker_id, worker=num_workers, seed=seed)
        dataset = iter(dataset)
        fetcher = stream_fetcher

    while watchdog.is_alive():
        try:
            r = index_queue.get(timeout=GLOBAL_TIMEOUT)
        except queue.Empty:
            continue
        if r is None:
            assert done_event.is_set() or iteration_end
            break
        elif done_event.is_set() or iteration_end:
            continue

        idx, place_holder = r
        try:
            if data_monitor:
                data, pid, dataset_time, transform_time, collate_time = fetcher(
                    dataset,
                    place_holder,
                    transform,
                    collate,
                    parallel_stream,
                    batch_size,
                )

                if idx < monitor_num_workers.value:
                    if monitor_workers[idx] == 0:
                        monitor_workers[idx] = pid
                worker_idx = idx % monitor_num_workers.value
                if (
                    pid == monitor_workers[worker_idx]
                    and idx % data_monitor_frequency == 0
                ):
                    print(