torch.cuda.streams.Stream in multiprocess context causes error: can't pickle Stream objects

username: Alex_Luya
cooked: <p>Hello,I added data prefetching by using cuda stream like this:</p>
<pre><code class="lang-auto">class data_prefetcher():
    def __init__(self, loader):
        self.loader = iter(loader)
        self.stream = torch.cuda.Stream()
        self.mean = torch.tensor([0.485 * 255, 0.456 * 255, 0.406 * 255]).cuda().view(1,3,1,1)
        self.std = torch.tensor([0.229 * 255, 0.224 * 255, 0.225 * 255]).cuda().view(1,3,1,1)
        self.preload()

    def preload(self):
        try:
            self.next_input, self.next_target = next(self.loader)
        except StopIteration:
            self.next_input = None
            self.next_target = None
            return

        with torch.cuda.stream(self.stream):
            self.next_input = self.next_input.cuda(non_blocking=True)
            self.next_target = self.next_target.cuda(non_blocking=True)
          
            self.next_input = self.next_input.float()
            self.next_input = self.next_input.sub_(self.mean).div_(self.std)

    def next(self):
        torch.cuda.current_stream().wait_stream(self.stream)
        input = self.next_input
        target = self.next_target
        if input is not None:
            input.record_stream(torch.cuda.current_stream())
        if target is not None:
            target.record_stream(torch.cuda.current_stream())
        self.preload()
        return input, target
</code></pre>
<p>(come from <a href="https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py#L256" rel="nofollow noopener">https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py#L256</a>)</p>
<p>to training logic,it works fine in single GPU training,but when moving to multiprocess context,I got error like this:</p>
<pre><code class="lang-auto">  File "/usr/lib/python3.6/multiprocessing/process.py", line 105, in start
    self._popen = self._Popen(self)
  File "/usr/lib/python3.6/multiprocessing/context.py", line 223, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/usr/lib/python3.6/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "/usr/lib/python3.6/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/usr/lib/python3.6/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/usr/lib/python3.6/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/usr/lib/python3.6/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
TypeError: can't pickle Stream objects
</code></pre>
<p>After debuging,I found <strong>torch.cuda.streams.Stream</strong> triggered this exception,Question are:</p>
<p>1,Isn’t it possible to use cuda stream in torch.multiprocess context?<br>
2,If not,any examples?</p>

username: mrshenli
cooked: <p>Hi <a class="mention" href="/u/alex_luya">@Alex_Luya</a></p>
<p>If all you need is syncing streams across processes, you can use the <a href="https://pytorch.org/docs/stable/cuda.html#torch.cuda.Event.ipc_handle" rel="nofollow noopener">ipc_handle()</a> API to pass CUDA events across processes. See the example in the <a href="https://github.com/pytorch/pytorch/blob/f314d9a0774062a20015ae522d33eadd45293328/test/test_multiprocessing.py#L590-L623" rel="nofollow noopener">test</a>.</p>