rework dry_run handling in mnist, mnist_hogwild

Author: mattip

mnist/main.py

@@ -47,6 +47,8 @@ def train(args, model, device, train_loader, optimizer, epoch):
             print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                 epoch, batch_idx * len(data), len(train_loader.dataset),
                 100. * batch_idx / len(train_loader), loss.item()))
+            if args.dry_run:
+                break
 
 
 def test(model, device, test_loader):
@@ -113,26 +115,6 @@ def main():
                        transform=transform)
     dataset2 = datasets.MNIST('../data', train=False,
                        transform=transform)
-    if args.dry_run:
-        from torch.utils.data.sampler import Sampler
-
-        class DryRunSampler(Sampler):
-            r"""Return only two datum from the set of data
-            """
-
-            def __init__(self, data_source):
-                self.data_source = data_source
-
-            def __iter__(self):
-                return iter(range(2))
-
-            def __len__(self):
-                return 2
-
-
-        kwargs['sampler'] = DryRunSampler(dataset1)
-        kwargs['shuffle'] = False
-
     train_loader = torch.utils.data.DataLoader(dataset1,**kwargs)
     test_loader = torch.utils.data.DataLoader(dataset2, **kwargs)
 

mnist_hogwild/main.py

@@ -51,20 +51,6 @@ def forward(self, x):
         return F.log_softmax(x, dim=1)
 
 
-class DryRunSampler(Sampler):
-    r"""Return only two datum from the set of data
-    """
-
-    def __init__(self, data_source):
-        self.data_source = data_source
-
-    def __iter__(self):
-        return iter(range(2))
-
-    def __len__(self):
-        return 2
-
-
 if __name__ == '__main__':
     args = parser.parse_args()
 
@@ -85,12 +71,6 @@ def __len__(self):
                        'pin_memory': True,
                       })
 
-    if args.dry_run:
-
-        kwargs['sampler'] = DryRunSampler(dataset1)
-        kwargs['shuffle'] = False
-
-
     torch.manual_seed(args.seed)
     mp.set_start_method('spawn')
 

mnist_hogwild/train.py

@@ -35,6 +35,8 @@ def train_epoch(epoch, args, model, device, data_loader, optimizer):
             print('{}\tTrain Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                 pid, epoch, batch_idx * len(data), len(data_loader.dataset),
                 100. * batch_idx / len(data_loader), loss.item()))
+            if args.dry_run:
+                break
 
 
 def test_epoch(model, device, data_loader):

run_python_examples.sh

@@ -172,7 +172,7 @@ function clean() {
 }
 
 function run_all() {
-  #cpp
+  # cpp
   dcgan
   # distributed
   fast_neural_style