modified the hogwild example to perform testing outside the different…

… training processes (#426)

mnist_hogwild/main.py

@@ -5,7 +5,7 @@
 import torch.nn.functional as F
 import torch.multiprocessing as mp
 
-from train import train
+from train import train, test
 
 # Training settings
 parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
@@ -55,7 +55,13 @@ def forward(self, x):
     processes = []
     for rank in range(args.num_processes):
         p = mp.Process(target=train, args=(rank, args, model))
+        # 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)
+
+

mnist_hogwild/train.py

@@ -14,17 +14,22 @@ def train(rank, args, model):
                         transforms.Normalize((0.1307,), (0.3081,))
                     ])),
         batch_size=args.batch_size, shuffle=True, num_workers=1)
-    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)
 
     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)
-        test_epoch(model, test_loader)
+
+def test(args, model):
+    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)
+
+    test_epoch(model, test_loader)
 
 
 def train_epoch(epoch, args, model, data_loader, optimizer):