PyTorch DataParallel not using second GPU during Inference

sree_harsha: <p>Hello,</p>
<p>I’m trying to run inference on a MMSR model. The system has two 2080Ti GPUs and I’m running PyTorch 1.1.0 on Ubuntu 18.04 with CUDA 10.</p>
<p>I have wrapped the model around a <code>nn.DataParallel</code>, and have tried setting and ensured that the two GPUs are visible<br>
<code>os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1"</code></p>
<p><code>torch.cuda.device_count()</code> returned 2.</p>
<p>Can anyone help me out? The inference script just goes CUDA OOM and doesn’t use the second GPU.</p>
<p>Edit:<br>
Code Snippet</p>
<pre><code class="lang-auto">model.load_state_dict(torch.load(model_path), strict=True)
model.eval()  
model = nn.DataParallel(model, devices_ids=["0, 1"])                                                                 
model = model.to(device)
</code></pre>
<p>I’ve looked at all tutorials on the website and used it before, haven’t come across this behaviour. Not sure what I’m doing wrong.</p>
<p>For some more context, I’m trying to load one batch at a time of 5 pretty large images (no way around this) I have a vague feeling I can’t use this approach, so any suggestions would be appreicated.</p>
<p>Thank you so much.</p>

ptrblck: <p>Since your device runs out of memory, you would need to reduce the batch size.<br>
Also, for inference you could warp the code in a <code>with torch.no_grad()</code> block to save some memory.</p>

sree_harsha: <p>Thank you for the reply.<br>
I am utilizing torch.no_grad() and I have to load 5 images per run of the algorithm.</p>
<p>The question was more around why PyTorch is not utilizing DataParallel as intended.</p>
<p>This is for (1) batch.</p>
<p>This is more a case of the model itself being large, would it require explicit ModelParallel sort of optimization?</p>

ptrblck: <aside class="quote no-group" data-username="sree_harsha" data-post="3" data-topic="74817">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sree_harsha/48/14426_2.png" class="avatar"> sree_harsha:</div>
<blockquote>
<p>This is for (1) batch.</p>
</blockquote>
</aside>
<p>Does this mean you are using a batch size of 1?<br>
If so, <code>nn.DataParallel</code> cannot be used, as the batch will be chunked in the batch dimension (dim0) and each chunk will be scattered to the corresponding device.</p>
<aside class="quote no-group" data-username="sree_harsha" data-post="3" data-topic="74817">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sree_harsha/48/14426_2.png" class="avatar"> sree_harsha:</div>
<blockquote>
<p>This is more a case of the model itself being large, would it require explicit ModelParallel sort of optimization?</p>
</blockquote>
</aside>
<p>If the model is too large for a single GPU, you could use model parallel, yes.</p>
<p>However, I’m currently unsure which use cases work.<br>
What’s the largest batch size for a single GPU and what batch sizes are you using for the data parallel approach?</p>

sree_harsha: <p>I’m terribly sorry for not noticing this message earlier.</p>
<p>The largest batch size is 1.<br>
But 1 batch contains 5 images which are of a certain resolution say [1920, 1080]<br>
So the shape of one batch is [1, 5, 3, 1920, 1080].</p>
<p>Downsampling the image in this situation is not possible due to objective, and the per forward pass all 5 images are required. I think the only way to address this problem seems to be to increase the GPU VRAM, but I am already tried using a Tesla V100.</p>
<p>I was wondering, if considering there are two GPUs with say 16 GB + 16 GB, is there a way to distribute inference across both of them (via ModelParallel) or otherwise, to potentially process them this way?</p>
<p>And is there any other way to deal with this? I was interested if this was possible.</p>

ptrblck: <aside class="quote no-group" data-username="sree_harsha" data-post="5" data-topic="74817">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sree_harsha/48/14426_2.png" class="avatar"> sree_harsha:</div>
<blockquote>
<p>I was wondering, if considering there are two GPUs with say 16 GB + 16 GB, is there a way to distribute inference across both of them (via ModelParallel) or otherwise, to potentially process them this way?</p>
</blockquote>
</aside>
<p>Yes, this would be possible.<a href="https://discuss.pytorch.org/t/split-single-model-in-multiple-gpus/13239/2">Here</a> is a simple example of model sharding.<br>
Basically you can push submodules to specific devices and would have to make sure to push the activation in the <code>forward</code> method to the right device.</p>
<p>Let me know, if you get stuck or need more information.</p>

sree_harsha: <p>The example looks great, I’ll try it out and get back to you!</p>
<p>Thank you <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>