diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/pytorch_lightning/examples/__init__.py b/pytorch_lightning/examples/__init__.py new file mode 100644 index 0000000000..6743d7f979 --- /dev/null +++ b/pytorch_lightning/examples/__init__.py @@ -0,0 +1 @@ +from .new_project_templates.lightning_module_template import LightningTemplateModel \ No newline at end of file diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000000..5dec59a5b8 --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,92 @@ +import pytest +from pytorch_lightning import Trainer +from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel +from argparse import Namespace +from test_tube import Experiment +import os + + +def get_model(): + root_dir = os.path.dirname(os.path.realpath(__file__)) + hparams = Namespace(**{'drop_prob': 0.2, + 'batch_size': 32, + 'in_features': 28*28, + 'learning_rate': 0.001*8, + 'optimizer_name': 'adam', + 'data_root': os.path.join(root_dir, 'mnist'), + 'out_features': 10, + 'hidden_dim': 1000}) + model = LightningTemplateModel(hparams) + + return model + +def get_exp(): + exp = Experiment(debug=True) + return exp + +def test_cpu_model(): + model = get_model() + + trainer = Trainer( + experiment=get_exp(), + max_nb_epochs=1, + train_percent_check=0.4, + val_percent_check=0.4 + ) + + result = trainer.fit(model) + + assert result == 1 + + +def test_single_gpu_model(): + model = get_model() + + trainer = Trainer( + experiment=get_exp(), + max_nb_epochs=1, + train_percent_check=0.4, + val_percent_check=0.4, + gpus=[0] + ) + + result = trainer.fit(model) + + assert result == 1 + + +def test_multi_gpu_model_dp(): + model = get_model() + + trainer = Trainer( + experiment=get_exp(), + max_nb_epochs=1, + train_percent_check=0.4, + val_percent_check=0.4, + gpus=[0, 1] + ) + + result = trainer.fit(model) + + assert result == 1 + + +def test_multi_gpu_model_ddp(): + model = get_model() + + trainer = Trainer( + experiment=get_exp(), + max_nb_epochs=1, + train_percent_check=0.4, + val_percent_check=0.4, + gpus=[0, 1], + distributed_backend='ddp' + ) + + result = trainer.fit(model) + + assert result == 1 + + +if __name__ == '__main__': + pytest.main([__file__])