You can use SimpleQueue in torch.multiprocessing to do that. E.g., you can create a queue between the host process and each subprocess, and use the queue to pass input data and collect output. The test below can serve as an example:

  
      github.com
  
  
    pytorch/pytorch/blob/0c5cd8c2b9cdf473e30bbb1b49ca80ed442813df/test/test_multiprocessing.py#L577-L600 1

@unittest.skipIf(NO_MULTIPROCESSING_SPAWN, "Disabled for environments that \
                 don't support multiprocessing with spawn start method")
@unittest.skipIf(not TEST_CUDA_IPC, 'CUDA IPC not available')
def test_event_multiprocess(self):
    event = torch.cuda.Event(enable_timing=False, interprocess=True)
    self.assertTrue(event.query())

    ctx = mp.get_context('spawn')
    p2c = ctx.SimpleQueue()
    c2p = ctx.SimpleQueue()
    p = ctx.Process(
        target=TestMultiprocessing._test_event_multiprocess_child,
        args=(event, p2c, c2p))
    p.start()

    c2p.get()  # wait for until child process is ready
    torch.cuda._sleep(50000000)  # spin for about 50 ms
    event.record()
    p2c.put(0)  # notify child event is recorded



  This file has been truncated. show original