Dataparallel with custom model

PhysicsIsFun: <p>Greetings,</p>
<p>I don’t quite get how to use the Dataprallel wrapper to use multiple GPUs for my custom model.</p>
<p>In my particular case, I wrote my model with an <code>evaluate</code> member function that already uses the device. In my torch framework, all my train routines expect the models to have this <code>evaluate</code> function.</p>
<p>My model is a custom Resnet model, built on ResNet18 from TorchVision.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch.nn as nn
import torchvision
import torch

class ResNet(nn.Module):
    def __init__(self, model, device):
        super(ResNet, self).__init__()
        self.resnet = model
        self.device = device
        if not isinstance(model, torchvision.models.ResNet):
            raise ValueError("The given model is not an instance of resnet.")

    def forward(self, features):
        return self.resnet.forward(features)
    
    def evaluate(self, data_loader):
        self.eval()
        loss = 0
        correct = 0
        criterion = nn.CrossEntropyLoss(reduction="sum")
        with torch.no_grad():
            for data, target in data_loader:
                data, target = data.to(self.device), target.to(self.device)
                output = self.forward(data)
                loss += criterion(output, target).data.item()               # sum up batch loss
                correct += (target == torch.argmax(output, dim=1)).cpu().sum() # does cpu make sense?
    
        accu = 100. * correct / len(data_loader.dataset)
        
        return (loss, accu)
    

model = torchvision.models.resnet18()  ## get predefined resnet.

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')  

res_model = ResNet(model, device)  ## instantiate custom model

res_model_wrapped = nn.DataParallel(res_model,device_ids = [0,3])  ## wrap  in DataParallel for GPUS 0 and 3
res_model_wrapped.to(device)


### call train routine on `res_model_wrapped.module`
</code></pre>
<p>The code runs like this, but only one of the GPU is working, the one with index 0.<br>
Any idea how I can make this work?</p>

Eduardo_Lawson: <p>When you use it as this:</p>
<pre><code class="lang-auto">torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')  
</code></pre>
<p>You are specifying that you will be using the one with index 0.<br>
If you want to use all the available GPUs:</p>
<pre><code class="lang-auto">device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
</code></pre>
<p>But if you want to use with specific GPUs:</p>
<pre><code class="lang-auto">device = torch.device("cuda:1,3" if torch.cuda.is_available() else "cpu")
</code></pre>

PhysicsIsFun: <p>Using<br>
<code>device = torch.device("cuda:0,3" if torch.cuda.is_available() else "cpu")</code><br>
gives me<br>
<code>RuntimeError: Invalid device string: 'cuda:0,3' </code></p>

ptrblck: <aside class="quote no-group" data-username="PhysicsIsFun" data-post="1" data-topic="204779">
<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/p/f6c823/48.png" class="avatar"> PhysicsIsFun:</div>
<blockquote>
<pre><code class="lang-auto">### call train routine on `res_model_wrapped.module`
</code></pre>
</blockquote>
</aside>
<p>This is wrong as it will skip the <code>nn.DataParallel</code> wrapper and call the internal module directly. <code>nn.DataParallel</code> will chunk the input data and will use the <code>forward</code> method of your model for each input chunk on the corresponding device.</p>

PhysicsIsFun: <p>If I call the train routine simply on <code>res_model_wrapped</code>,<br>
I get the error<br>
<code>AttributeError: 'DataParallel' object has no attribute 'evaluate'</code>, because the train routine calls the <code>evaluate</code> function.</p>
<p>Should I call <code>res_model_wrapped.module.evaluate()</code> at these points? Will it be a problem that the evaluate function uses <code>device</code>?</p>

ptrblck: <p>You would need to call custom methods from the <code>forward</code> function without hard-coding any device, since the <code>forward</code> is used in <code>nn.DataParallel</code> for the actual data parallel application. Using the internal <code>.module</code> will skip <code>nn.DataParallel</code> and you will use the single-device module again.</p>

PhysicsIsFun: <p>So my train routine should not call any member functions of the module wrapped in <code>nn.DataParallel</code> that explicitly uses <code>device</code>?</p>