DataParallel with non-torch objects that requires moving to device

roaiyadgar: <p>Hi,<br>
I’m trying to use DataParallel on my Module, but one of the attributes of the module (<code>nufft_plan</code>) isn’t a torch object but it still dependent on the module’s device:</p>
<pre data-code-wrap="python"><code class="lang-python">class MyModule(torch.nn.Module):
    def __init__(self,data,params):
        super(MyModule,self).__init__()
        self.nufft_plan = NufftPlan(params,data.device)
        self.data= torch.nn.Parameter(data)
        self.params = params

    def forward(self,input):
        output = some_computation(self.data,self.nufft_plan,input)
        return output
    
    def to(self,device):
        module = super().to(device)
        module.nufft_plan = NufftPlan(self.params,device)

        return module
</code></pre>
<p>When I use <code>DataParallel</code> on <code>MyModule</code> it doesn’t invoke the <code>.to</code> method, so I can’t move <code>nufft_plan</code> into each GPU device. Is there any other way I can move <code>nufft_plan</code> into the right device when using <code>DataParallel</code>?</p>

ptrblck: <p>You can move it to the right device inside the <code>forward</code> method using the <code>.device</code> attribute of any other properly registered parameter or the input.</p>

roaiyadgar: <p>Hi,<br>
This is what I will probably end up doing, but I’d rather avoid moving it in forward, since this object doesn’t actually allow me to move it to a different device and I just overwrite it with a new object on the desired device. It also doesn’t allow me to check what it’s current device.<br>
If I do the moving in forward I need to check if the object is on the right device or not (so<br>
I must save its current device) and move it accordingly.<br>
Is there a more “elegant” way of doing it? something like using a custom function when wrapping the module with DataParallel?</p>
<p>Thank you very much <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

ptrblck: <p>The more elegant and faster approach would be to use <code>DistributedDataParallel</code> which avoids model copies performed in <code>nn.DataParallel</code> and allows you to use a single process per device. Since each process moves its model to the corresponding device only, the repeated calls in the <code>forward</code> wouldn’t be needed anymore.</p>

roaiyadgar: <p>I see, I did use <code>DistributedDataParallel</code>(on something else) but the computation I’m trying to do here could be relatively short and then the time it takes to start the processes can actually slow it down. I guess I could use both DataParallel/DistrubtedDataParallel depending on the size of the dataset (I’m not trying to train some model but just perform some computation over the dataset).</p>
<p>Thank you very much for your help <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

Nishant_Madan: <p>Hi,</p>
<p>Model Parallelism(TensorParallelism or PipelineParallelism) not DataParallel or DataDistributedParallel.</p>
<p>I am a master student and I am also trying to do the related task which is performing the Distributed Inference on Edge devices but I am facing some challenges while doing the Pytorch RPC implementation to connect to multiple devices using their IP addresses and port. For now I have 2 Nvidia Jetson Nanos and their IP addresses and I am not sure what is the correct way of serving the NN layers on both the edge devices simultaneously.</p>
<p>Any help would be great. Thank you in advance.</p>