DDP and Spectral Normalization

ClementPla: <p>Hi,<br>
I’m having trouble with spectral normalization when training a distributed model.<br>
The code works fine on a single GPU, but as soon as I move to multiples GPU, the following message raises:</p>
<pre><code class="lang-auto">File "/usagers/clpla/NNTools/nntools/experiment/experiment.py", line 345, in start
    mp.spawn(self._start_process, nprocs=self.world_size, join=True)
  File "/usagers/clpla/.conda/envs/torch18/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 230, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/usagers/clpla/.conda/envs/torch18/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 179, in start_processes
    process.start()
  File "/usagers/clpla/.conda/envs/torch18/lib/python3.8/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/usagers/clpla/.conda/envs/torch18/lib/python3.8/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "/usagers/clpla/.conda/envs/torch18/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/usagers/clpla/.conda/envs/torch18/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/usagers/clpla/.conda/envs/torch18/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/usagers/clpla/.conda/envs/torch18/lib/python3.8/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
  File "/usagers/clpla/.conda/envs/torch18/lib/python3.8/site-packages/torch/nn/utils/parametrize.py", line 285, in getstate
    raise RuntimeError(
RuntimeError: Serialization of parametrized modules is only supported through state_dict(). See:
https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-a-general-checkpoint-for-inference-and-or-resuming-training
Exception ignored in: &lt;function Connection.__del__ at 0x7fea8bff7ca0&gt;

</code></pre>
<p>I simply use Pytorch’s implementation of spectral normalization:</p>
<pre><code class="lang-auto">model = my_model()
for m in model.modules():
    if type(m) == nn.Conv2d or type(m) == nn.Linear:
        nn.utils.parametrizations.spectral_norm(m)
</code></pre>
<p>Any idea on how to fix this issue? I’m using Pytorch 1.10.2.<br>
Thanks!</p>

pritamdamania87: <p>It looks like you are first initializing your module/spectral_norm and then trying to pass it through <code>mp.spawn</code> which is trying to pickle it and failing. I’d suggest running <code>mp.spawn</code> first and then within each process initialize your model. You can find simple examples of this here: <a href="https://pytorch.org/tutorials/intermediate/dist_tuto.html" class="inline-onebox" rel="noopener nofollow ugc">Writing Distributed Applications with PyTorch — PyTorch Tutorials 1.10.1+cu102 documentation</a></p>

ClementPla: <p>Hi, thanks for your answer and sorry for providing a late feedback,<br>
I was trying to see how to adapt my code for this specific case, my pipeline wasn’t initially thought with this problem in mind.<br>
Initializing the model after the call to mp.spawn did solve the problem, thanks!</p>