speed-up testing (#504)

* extend CI timeout

* add short MNIST

* lower dataset and stop thr

* refactor imports

* formatting

* early stop

* play params

* play params

* minor refactoring

# Conflicts:
#	pytorch_lightning/testing/__init__.py
#	pytorch_lightning/testing/lm_test_module.py
#	pytorch_lightning/testing/lm_test_module_base.py
#	pytorch_lightning/testing/lm_test_module_mixins.py
#	pytorch_lightning/testing/model.py
#	pytorch_lightning/testing/model_base.py
#	pytorch_lightning/testing/model_mixins.py
#	pytorch_lightning/testing/test_module.py
#	pytorch_lightning/testing/test_module_base.py
#	pytorch_lightning/testing/test_module_mixins.py

* typo

Co-Authored-By: Ir1dXD <sirius.caffrey@gmail.com>

* Revert "refactor imports"

This reverts commit b86aee92

* update imports

pytorch_lightning/testing/__init__.py

@@ -1,6 +1,6 @@
-from .test_module import LightningTestModel
-from .test_module_base import LightningTestModelBase
-from .test_module_mixins import (
+from .model import LightningTestModel
+from .model_base import LightningTestModelBase
+from .model_mixins import (
     LightningValidationStepMixin,
     LightningValidationMixin,
     LightningValidationStepMultipleDataloadersMixin,

pytorch_lightning/testing/model.py

@@ -1,7 +1,7 @@
 import torch
 
-from .test_module_base import LightningTestModelBase
-from .test_module_mixins import LightningValidationMixin, LightningTestMixin
+from .model_base import LightningTestModelBase
+from .model_mixins import LightningValidationMixin, LightningTestMixin
 
 
 class LightningTestModel(LightningValidationMixin, LightningTestMixin, LightningTestModelBase):

pytorch_lightning/testing/model_base.py

@@ -19,6 +19,22 @@ from pytorch_lightning import data_loader
 from pytorch_lightning.core.lightning import LightningModule
 
 
+class TestingMNIST(MNIST):
+
+    def __init__(self, root, train=True, transform=None, target_transform=None,
+                 download=False, num_samples=8000):
+        super(TestingMNIST, self).__init__(
+            root,
+            train=train,
+            transform=transform,
+            target_transform=target_transform,
+            download=download
+        )
+        # take just a subset of MNIST dataset
+        self.data = self.data[:num_samples]
+        self.targets = self.targets[:num_samples]
+
+
 class LightningTestModelBase(LightningModule):
     """
     Base LightningModule for testing. Implements only the required
@@ -137,8 +153,8 @@ class LightningTestModelBase(LightningModule):
         # init data generators
         transform = transforms.Compose([transforms.ToTensor(),
                                         transforms.Normalize((0.5,), (1.0,))])
-        dataset = MNIST(root=self.hparams.data_root, train=train,
-                        transform=transform, download=True)
+        dataset = TestingMNIST(root=self.hparams.data_root, train=train,
+                               transform=transform, download=True, num_samples=2000)
 
         # when using multi-node we need to add the datasampler
         train_sampler = None