Making a wrapper around nn.DataParallel to access module attributes is safe?

username: hilman-dayo
cooked: <p>Hello everyone.<br>
I am using pytorch <code>nn.DataParallel</code>, but I cannot access any variables within may model with it (I will get AttributeError).</p>
<p>So, I make a wrapper like this:</p>
<pre><code class="lang-auto">class DataParallelWrapper(nn.DataParallel):
    def __getattr__(self, name):
        try:
            return super().__getattr__(name)
        except AttributeError:
            return getattr(self.module, name)
</code></pre>
<p>Being completely zero in how <code>nn.DataParallel</code> is being implemented, it raises a question.<br>
Is this a safe approach?</p>
<p>If it matters, I am using PyTorch 1.1 (need to update code before upgrading to new shiny version).</p>

username: mrshenli
cooked: <p><code>DataParallel</code> does have a <code>self.module</code> variable:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/df8d6eeb19423848b20cd727bc4a728337b73829/torch/nn/parallel/data_parallel.py#L131" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/df8d6eeb19423848b20cd727bc4a728337b73829/torch/nn/parallel/data_parallel.py#L131" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/df8d6eeb19423848b20cd727bc4a728337b73829/torch/nn/parallel/data_parallel.py#L131</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="121" style="counter-reset: li-counter 120 ;">
<li>        self.module = module</li>
<li>        self.device_ids = []</li>
<li>        return</li>
<li>
</li>
<li>    if device_ids is None:</li>
<li>        device_ids = list(range(torch.cuda.device_count()))</li>
<li>    if output_device is None:</li>
<li>        output_device = device_ids[0]</li>
<li>
</li>
<li>    self.dim = dim</li>
<li class="selected">    self.module = module</li>
<li>    self.device_ids = list(map(lambda x: _get_device_index(x, True), device_ids))</li>
<li>    self.output_device = _get_device_index(output_device, True)</li>
<li>    self.src_device_obj = torch.device("cuda:{}".format(self.device_ids[0]))</li>
<li>
</li>
<li>    _check_balance(self.device_ids)</li>
<li>
</li>
<li>    if len(self.device_ids) == 1:</li>
<li>        self.module.cuda(device_ids[0])</li>
<li>
</li>
<li>def forward(self, *inputs, **kwargs):</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>The following code works for me:</p>
<pre><code class="lang-python">import torch.nn as nn
import torch
from torch.nn.parallel import DataParallel

class DataParallelWrapper(DataParallel):
    def __init__(self, module):
        super().__init__(module)


    def __getattr__(self, name):
        try:
            return super().__getattr__(name)
        except AttributeError:
            return getattr(self.module, name)


dpw = DataParallelWrapper(nn.Linear(2, 2))
print(getattr(dpw, "forward"))
print(getattr(dpw, "weight"))
</code></pre>

username: hilman-dayo
cooked: <p>I am fully aware of the ability of accessing my network through <code>nn.DataParallel.module</code>. However, I need to make sure that my code is executable (no AttributeError when I am trying to access my defined-variable in my network) whether the user is using single GPU or multiple GPUs. Hence, the <code>DataPrallelWrapper</code>.</p>
<p>The only problem is I do not know whether this is a good approach or not (will my approach invite any data race or async problem?).</p>

username: mrshenli
cooked: <aside class="quote no-group" data-username="hilman-dayo" data-post="3" data-topic="79124">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/ba9def/48.png" class="avatar"> hilman-dayo:</div>
<blockquote>
<p>will my approach invite any data race or async problem?</p>
</blockquote>
</aside>
<p>This should be fine I think. The main logic of <code>DataParallel</code> is implemented in the <code>forward()</code> function, it replicates the given module to available devices, scatters input, launches one thread per device to process the scattered input, and then joins threads and gathers the outputs. So as long as you are not trying to modify module parameters or gradients concurrently when executing <code>DataParallel.forward()</code>, I don’t see an issue here.</p>