Add Cuda support to mnist_hogwild (pytorch#508)

Author: VitalyFedyunin

mnist_hogwild/main.py

@@ -25,6 +25,8 @@
                     help='how many batches to wait before logging training status')
 parser.add_argument('--num-processes', type=int, default=2, metavar='N',
                     help='how many training processes to use (default: 2)')
+parser.add_argument('--cuda', action='store_true', default=False,
+                    help='enables CUDA training')
 
 class Net(nn.Module):
     def __init__(self):
@@ -47,21 +49,24 @@ def forward(self, x):
 if __name__ == '__main__':
     args = parser.parse_args()
 
+    use_cuda = args.cuda and torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+    dataloader_kwargs = {'pin_memory': True} if use_cuda else {}
+
     torch.manual_seed(args.seed)
+    mp.set_start_method('spawn')
 
-    model = Net()
+    model = Net().to(device)
     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))
+        p = mp.Process(target=train, args=(rank, args, model, device, dataloader_kwargs))
         # We first train the model across `num_processes` processes
         p.start()
         processes.append(p)
     for p in processes:
         p.join()
 
     # Once training is complete, we can test the model
-    test(args, model)
-
-
+    test(args, model, device, dataloader_kwargs)

mnist_hogwild/train.py

@@ -4,7 +4,8 @@
 import torch.nn.functional as F
 from torchvision import datasets, transforms
 
-def train(rank, args, model):
+
+def train(rank, args, model, device, dataloader_kwargs):
     torch.manual_seed(args.seed + rank)
 
     train_loader = torch.utils.data.DataLoader(
@@ -13,32 +14,35 @@ def train(rank, args, model):
                         transforms.ToTensor(),
                         transforms.Normalize((0.1307,), (0.3081,))
                     ])),
-        batch_size=args.batch_size, shuffle=True, num_workers=1)
+        batch_size=args.batch_size, shuffle=True, num_workers=1,
+        **dataloader_kwargs)
 
     optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
     for epoch in range(1, args.epochs + 1):
-        train_epoch(epoch, args, model, train_loader, optimizer)
+        train_epoch(epoch, args, model, device, train_loader, optimizer)
+
 
-def test(args, model):
+def test(args, model, device, dataloader_kwargs):
     torch.manual_seed(args.seed)
 
     test_loader = torch.utils.data.DataLoader(
         datasets.MNIST('../data', train=False, transform=transforms.Compose([
             transforms.ToTensor(),
             transforms.Normalize((0.1307,), (0.3081,))
         ])),
-        batch_size=args.batch_size, shuffle=True, num_workers=1)
+        batch_size=args.batch_size, shuffle=True, num_workers=1,
+        **dataloader_kwargs)
 
-    test_epoch(model, test_loader)
+    test_epoch(model, device, test_loader, device)
 
 
-def train_epoch(epoch, args, model, data_loader, optimizer):
+def train_epoch(epoch, args, model, device, data_loader, optimizer):
     model.train()
     pid = os.getpid()
     for batch_idx, (data, target) in enumerate(data_loader):
         optimizer.zero_grad()
-        output = model(data)
-        loss = F.nll_loss(output, target)
+        output = model(data.to(device))
+        loss = F.nll_loss(output, target.to(device))
         loss.backward()
         optimizer.step()
         if batch_idx % args.log_interval == 0:
@@ -47,16 +51,16 @@ def train_epoch(epoch, args, model, data_loader, optimizer):
                 100. * batch_idx / len(data_loader), loss.item()))
 
 
-def test_epoch(model, data_loader):
+def test_epoch(model, device, data_loader):
     model.eval()
     test_loss = 0
     correct = 0
     with torch.no_grad():
         for data, target in data_loader:
-            output = model(data)
-            test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
+            output = model(data.to(device))
+            test_loss += F.nll_loss(output, target.to(device), reduction='sum').item() # sum up batch loss
             pred = output.max(1)[1] # get the index of the max log-probability
-            correct += pred.eq(target).sum().item()
+            correct += pred.eq(target.to(device)).sum().item()
 
     test_loss /= len(data_loader.dataset)
     print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(