clean test code

hapi/test/dist_mnist.py

@@ -29,92 +29,22 @@ from paddle.io import BatchSampler, DataLoader
 
 from hapi.model import Model, Input, set_device
 from hapi.loss import Loss, CrossEntropy
+from hapi.vision.models import LeNet
 from hapi.metrics import Accuracy
 from hapi.callbacks import ProgBarLogger
-from hapi.datasets import MNIST as MnistDataset
-
-
-class SimpleImgConvPool(fluid.dygraph.Layer):
-    def __init__(self,
-                 num_channels,
-                 num_filters,
-                 filter_size,
-                 pool_size,
-                 pool_stride,
-                 pool_padding=0,
-                 pool_type='max',
-                 global_pooling=False,
-                 conv_stride=1,
-                 conv_padding=0,
-                 conv_dilation=1,
-                 conv_groups=None,
-                 act=None,
-                 use_cudnn=False,
-                 param_attr=None,
-                 bias_attr=None):
-        super(SimpleImgConvPool, self).__init__('SimpleConv')
-
-        self._conv2d = Conv2D(
-            num_channels=num_channels,
-            num_filters=num_filters,
-            filter_size=filter_size,
-            stride=conv_stride,
-            padding=conv_padding,
-            dilation=conv_dilation,
-            groups=conv_groups,
-            param_attr=None,
-            bias_attr=None,
-            use_cudnn=use_cudnn)
-
-        self._pool2d = Pool2D(
-            pool_size=pool_size,
-            pool_type=pool_type,
-            pool_stride=pool_stride,
-            pool_padding=pool_padding,
-            global_pooling=global_pooling,
-            use_cudnn=use_cudnn)
-
-    def forward(self, inputs):
-        x = self._conv2d(inputs)
-        x = self._pool2d(x)
-        return x
-
-
-class MNIST(Model):
-    def __init__(self):
-        super(MNIST, self).__init__()
-        self._simple_img_conv_pool_1 = SimpleImgConvPool(
-            1, 20, 5, 2, 2, act="relu")
-
-        self._simple_img_conv_pool_2 = SimpleImgConvPool(
-            20, 50, 5, 2, 2, act="relu")
-
-        pool_2_shape = 50 * 4 * 4
-        SIZE = 10
-        scale = (2.0 / (pool_2_shape**2 * SIZE))**0.5
-        self._fc = Linear(
-            800,
-            10,
-            param_attr=fluid.param_attr.ParamAttr(
-                initializer=fluid.initializer.NormalInitializer(
-                    loc=0.0, scale=scale)),
-            act="softmax")
-
-    def forward(self, inputs):
-        inputs = fluid.layers.reshape(inputs, [-1, 1, 28, 28])
-        x = self._simple_img_conv_pool_1(inputs)
-        x = self._simple_img_conv_pool_2(x)
-        x = fluid.layers.flatten(x, axis=1)
-        x = self._fc(x)
-        return x
-
-
-class TestMnistDataset(MnistDataset):
-    def __init__(self):
-        super(TestMnistDataset, self).__init__(mode='test')
+from hapi.datasets import MNIST
+
+
+class MnistDataset(MNIST):
+    def __init__(self, mode, return_label=True):
+        super(MnistDataset, self).__init__(mode=mode)
+        self.return_label = return_label
 
     def __getitem__(self, idx):
-        return self.images[idx],
+        img = np.reshape(self.images[idx], [1, 28, 28])
+        if self.return_label:
+            return img, np.array(self.labels[idx]).astype('int64')
+        return img,
 
     def __len__(self):
         return len(self.images)
@@ -142,9 +72,9 @@ class TestModel(unittest.TestCase):
 
         train_dataset = MnistDataset(mode='train')
         val_dataset = MnistDataset(mode='test')
-        test_dataset = TestMnistDataset()
+        test_dataset = MnistDataset(mode='test', return_label=False)
 
-        model = MNIST()
+        model = LeNet()
         optim = fluid.optimizer.Momentum(
             learning_rate=0.01, momentum=.9, parameter_list=model.parameters())
         loss = CrossEntropy()

hapi/test/test_distributed.py

@@ -95,7 +95,6 @@ def start_local_trainers(cluster,
         print("trainer proc env:{}".format(current_env))
 
         cmd = "python -m coverage run --branch -p " + training_script
-        # cmd = [sys.executable, "-u", training_script] + training_script_args
 
         print("start trainer proc:{} env:{}".format(cmd, proc_env))