From 6ad542e2b6831ad2cc548e4aed361a98402cf1f7 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 24 Jul 2019 08:44:00 -0400
Subject: [PATCH] added gpu check for each gpu test

---
 tests/test_models.py | 93 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 93 insertions(+)

diff --git a/tests/test_models.py b/tests/test_models.py
index 5dec59a5..36bc87bf 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -3,6 +3,8 @@ 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 warnings
+import torch
 import os
 
 
@@ -36,10 +38,21 @@ def test_cpu_model():
 
     result = trainer.fit(model)
 
+    metrics = result.__tng_tqdm_dic
+    print(metrics)
+
     assert result == 1
 
 
 def test_single_gpu_model():
+    """
+    Make sure single GPU works (DP mode)
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_single_gpu_model cannot run. Rerun on a GPU node to run this test')
+        return
+
     model = get_model()
 
     trainer = Trainer(
@@ -56,6 +69,17 @@ def test_single_gpu_model():
 
 
 def test_multi_gpu_model_dp():
+    """
+    Make sure DP works
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
     model = get_model()
 
     trainer = Trainer(
@@ -72,6 +96,17 @@ def test_multi_gpu_model_dp():
 
 
 def test_multi_gpu_model_ddp():
+    """
+    Make sure DDP works
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
     model = get_model()
 
     trainer = Trainer(
@@ -88,5 +123,63 @@ def test_multi_gpu_model_ddp():
     assert result == 1
 
 
+def test_amp_gpu_ddp():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    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',
+        use_amp=True
+    )
+
+    result = trainer.fit(model)
+
+    assert result == 1
+
+
+def test_amp_gpu_dp():
+    """
+    Make sure DP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    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='dp',
+        use_amp=True
+    )
+
+    result = trainer.fit(model)
+
+    assert result == 1
+
+
 if __name__ == '__main__':
     pytest.main([__file__])