Is it possible to output an object instead of model when using multigpu?

tymons: <p>I have a Variational Autoencoder model with the following forward function:</p>
<pre><code class="lang-auto">    def forward(self, x):
        y = self.encoder(x)
        means = self.linear_means(y)
        log_var = self.linear_log_var(y)
        z = reparameterize(means, log_var)
        recon_x = self.decoder(z)
        return VaeOutput(output=recon_x, mean=means, log_var=log_var)
</code></pre>
<p>VaeOutput is an object:</p>
<pre><code class="lang-auto">@dataclass
class VaeOutput:
    output: torch.Tensor
    log_var: torch.Tensor
    mean: torch.Tensor
</code></pre>
<p>when I am training my model on a single GPU everything runs like a charm but whenever mu model is embedded with multi-gpus with nn.Dataparallel() i got an error about VaeOutput which is not serializable.</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "train.py", line 135, in &lt;module&gt;
    main()
  File "train.py", line 131, in main
    model = model_runner.train(model, train_loader, learning_config, val_loader, feature_config)
  File "/buzz-based-anomaly/utils/model_runner.py", line 375, in train
    return self._train(model, train_dataloader, train_config, self._train_step, self._val_step,
  File "/buzz-based-anomaly/utils/model_runner.py", line 455, in _train
    train_epoch_loss = train_step_func(model, train_dataloader, optimizer, experiment, epoch, log_interval)
  File "/buzz-based-anomaly/utils/model_runner.py", line 671, in _train_step
    model_output = model(batch)
  File "/root/miniconda3/envs/pytorch-cuda-env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/buzz-based-anomaly/utils/sm_data_parallel.py", line 10, in forward
    return self.model(*input_data)
  File "/root/miniconda3/envs/pytorch-cuda-env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/root/miniconda3/envs/pytorch-cuda-env/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 169, in forward
    return self.gather(outputs, self.output_device)
  File "/root/miniconda3/envs/pytorch-cuda-env/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 181, in gather
    return gather(outputs, output_device, dim=self.dim)
  File "/root/miniconda3/envs/pytorch-cuda-env/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 78, in gather
    res = gather_map(outputs)
  File "/root/miniconda3/envs/pytorch-cuda-env/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 73, in gather_map
    return type(out)(map(gather_map, zip(*outputs)))
TypeError: 'VaeOutput' object is not iterable
</code></pre>
<p>How should I approach that kind of problem?  Is it possible to use custom model output with multigpu support?</p>

eqy: <p>Would it be possible to wrap the model in another Class that contains a DataParallel instantiation of the model which then returns the object?</p>

wanchaol: <p>Thanks for posting <a class="mention" href="/u/tymons">@tymons</a>, I noticed that you are using nn.DataParallel to do mutigpu training, since we don’t recommend using <code>nn.DataParallel</code> right now, could you try using distributed data parallel (nn.DistributedDataParallel) and see if the issue still exist? You can see the tutorial here <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">Getting Started with Distributed Data Parallel — PyTorch Tutorials 1.10.0+cu102 documentation</a></p>

wanchaol: <p>actually DDP might also have this issue as internally it’s using Tuple[Tensor] to detect bucket and do optimizations, etc. Looking at your case, Could you try turning <code>VaeOutput</code> into Named Tuple, this way it still preserves its name and will be fully compatible with DDP</p>