FM / communication

/usr/local/lib64/python3.6/site-packages/torch/utils/data/communication/

/usr/local/lib64/python3.6/site-packages/torch/utils/data/communication

Name	Size	Mode	Actions
__pycache__/	-	0755	rm
eventloop.py	1536	0644	edit dl rm
iter.py	6279	0644	edit dl rm
messages.py	1171	0644	edit dl rm
protocol.py	5741	0644	edit dl rm
queue.py	1400	0644	edit dl rm
__init__.py	109	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/utils/data/communication/eventloop.py (1536B)

import torch
import threading
import pickle

from torch.utils.data import IterDataPipe, communication

def DataPipeToQueuesLoop(source_datapipe, req_queue, res_queue):
    if isinstance(source_datapipe, IterDataPipe):
        pipe_type = communication.iter
        protocol_type = communication.protocol.IterDataPipeQueueProtocolServer
    else:
        raise Exception('Only supports IterDataPipe, got', source_datapipe)
        # pipe_type = communication.map
        # protocol_type = communication.protocol.MapDataPipeQueueProtocolServer

torch.set_num_threads(1)
    for _ in pipe_type.DataPipeBehindQueues(source_datapipe, protocol_type(req_queue, res_queue), blocking_request_get=True):
        pass

def SpawnProcessForDataPipeline(multiprocessing_ctx, datapipe):
    req_queue = multiprocessing_ctx.Queue()
    res_queue = multiprocessing_ctx.Queue()
    process = multiprocessing_ctx.Process(
        target=DataPipeToQueuesLoop, args=(datapipe, req_queue, res_queue))
    return process, req_queue, res_queue

def SpawnThreadForDataPipeline(datapipe):
    req_queue = communication.queue.ThreadingQueue()
    res_queue = communication.queue.ThreadingQueue()

try:
        new_datapipe = pickle.loads(pickle.dumps(datapipe))
    except Exception as e:
        raise Exception('Unable to pickle DataPipe to make thread local copy', e)

process = threading.Thread(target=DataPipeToQueuesLoop, args=(
        new_datapipe, req_queue, res_queue), daemon=True)
    return process, req_queue, res_queue, new_datapipe

import torch import threading import pickle from torch.utils.data import IterDataPipe, communication def DataPipeToQueuesLoop(source_datapipe, req_queue, res_queue): if isinstance(source_datapipe, IterDataPipe): pipe_type = communication.iter protocol_type = communication.protocol.IterDataPipeQueueProtocolServer else: raise Exception('Only supports IterDataPipe, got', source_datapipe) # pipe_type = communication.map # protocol_type = communication.protocol.MapDataPipeQueueProtocolServer torch.set_num_threads(1) for _ in pipe_type.DataPipeBehindQueues(source_datapipe, protocol_type(req_queue, res_queue), blocking_request_get=True): pass def SpawnProcessForDataPipeline(multiprocessing_ctx, datapipe): req_queue = multiprocessing_ctx.Queue() res_queue = multiprocessing_ctx.Queue() process = multiprocessing_ctx.Process( target=DataPipeToQueuesLoop, args=(datapipe, req_queue, res_queue)) return process, req_queue, res_queue def SpawnThreadForDataPipeline(datapipe): req_queue = communication.queue.ThreadingQueue() res_queue = communication.queue.ThreadingQueue() try: new_datapipe = pickle.loads(pickle.dumps(datapipe)) except Exception as e: raise Exception('Unable to pickle DataPipe to make thread local copy', e) process = threading.Thread(target=DataPipeToQueuesLoop, args=( new_datapipe, req_queue, res_queue), daemon=True) return process, req_queue, res_queue, new_datapipe