MNIST example: clarify train/stats runs in code.

Author: lucasb-eyer

examples/MNIST/run.py

@@ -18,8 +18,8 @@ def main(params):
 
     for epoch in range(100):
         model.training()
-        train(train_set_x, train_set_y, model, optimiser, criterion, epoch, params['batch_size'])
-        train(train_set_x, train_set_y, model, optimiser, criterion, epoch, params['batch_size'], 'stat')
+        train(train_set_x, train_set_y, model, optimiser, criterion, epoch, params['batch_size'], 'train')
+        train(train_set_x, train_set_y, model, optimiser, criterion, epoch, params['batch_size'], 'stats')
 
         model.evaluate()
         validate(test_set_x, test_set_y, model, epoch, params['batch_size'])
@@ -29,4 +29,4 @@ def main(params):
     params = {}
     params['lr'] = 0.1
     params['batch_size'] = 64
-    main(params)
+    main(params)

examples/MNIST/train.py

@@ -3,8 +3,8 @@
 import theano as _th
 
 
-def train(dataset_x, dataset_y, model, optimiser, criterion, epoch, batch_size, mode=None):
-    progress = make_progressbar('Training', epoch, len(dataset_x))
+def train(dataset_x, dataset_y, model, optimiser, criterion, epoch, batch_size, mode='train'):
+    progress = make_progressbar('Training ({})'.format(mode), epoch, len(dataset_x))
     progress.start()
 
     shuffle = np.random.permutation(len(dataset_x))
@@ -17,12 +17,14 @@ def train(dataset_x, dataset_y, model, optimiser, criterion, epoch, batch_size,
             mini_batch_input[k] = dataset_x[shuffle[j * batch_size + k]]
             mini_batch_targets[k] = dataset_y[shuffle[j * batch_size + k]]
 
-        if mode is None:
+        if mode == 'train':
             model.zero_grad_parameters()
             model.accumulate_gradients(mini_batch_input, mini_batch_targets, criterion)
             optimiser.update_parameters(model)
-        else:
+        elif mode == 'stats':
             model.accumulate_statistics(mini_batch_input)
+        else:
+            assert False, "Mode should be either 'train' or 'stats'"
 
         progress.update(j * batch_size)