Hogwild on MultiGPU

bapi: <p>Hi,</p>
<p>I need to use multiple GPUs available in the machine in a way that each of the processes uses exactly one GPU. I modified the mnist_hogwild code <a href="https://github.com/pytorch/examples/blob/master/mnist_hogwild/main.py" rel="nofollow noopener">https://github.com/pytorch/examples/blob/master/mnist_hogwild/main.py</a> as the following:</p>
<pre><code class="lang-auto">dataloader_kwargs = {'pin_memory': True} if use_cuda else {}
    dcount = torch.cuda.device_count()
    devices = []
    model = Net()
    for i in range(dcount):
        devices.append(torch.device("cuda:"+str(i)))
    torch.manual_seed(args.seed)
    mp.set_start_method('spawn')

    # model = Net().to(device)
    for i in range(dcount):
        model.to(devices[i])
    model.share_memory() # gradients are allocated lazily, so they are not shared here

    processes = []
    for rank in range(args.num_processes):
        p = mp.Process(target=train, args=(rank, args, model, devices[int(rank%dcount)], dataloader_kwargs))
        # We first train the model across `num_processes` processes
        p.start()
        processes.append(p)
    for p in processes:
        p.join()

</code></pre>
<p>However, while running this code with num_processes = 2, as there are two GPUs in my machine, I can see only one of them engaged. Can you please suggest what exactly I need in the code here?</p>

auro: <p>Please review my version.<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/aurotripathy/menace/blob/master/mnist_multigpu_hogwild/main.py" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/aurotripathy/menace/blob/master/mnist_multigpu_hogwild/main.py" target="_blank" rel="nofollow noopener">aurotripathy/menace/blob/master/mnist_multigpu_hogwild/main.py</a></h4>
<pre><code class="lang-py">"""
Adding multi-gpu support to mnist w/hogwild
"""
from __future__ import print_function
import argparse
import torch
import torch.multiprocessing as mp
from model import Net
from train import train, test
from shared_optim import SharedAdam


# Training settings
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                    help='input batch size for training (default: 64)')
parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                    help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs', type=int, default=10, metavar='N',
                    help='number of epochs to train (default: 10)')
</code></pre>

  This file has been truncated. <a href="https://github.com/aurotripathy/menace/blob/master/mnist_multigpu_hogwild/main.py" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
<br>
I’m happy to fix issues, improve readability.</p>
<p>This really is derived from an RL implementation by <a class="mention" href="/u/dgriff">@dgriff</a> available <a href="https://github.com/dgriff777/rl_a3c_pytorch" rel="nofollow noopener">here</a></p>

pietern: <aside class="quote no-group" data-username="bapi" data-post="1" data-topic="44289">
<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/bapi/48/27077_2.png" class="avatar"> bapi:</div>
<blockquote>
<pre><code class="lang-auto">for i in range(dcount):
  model.to(devices[i])
</code></pre>
</blockquote>
</aside>
<p>This snippet will first move the model to device 0 and then to device 1. If you don’t explicitly move the model in the functions you’re running through multiprocessing, then you’ll have to make this dependent on the rank of the target process. As is, I assume you’re only using process 1.</p>

tengerye: <p>Hi, would you mind explaining why the shared optimizer is necessary please?</p>