Enable multi-gpu training in find_learning_rate.py (#2045)

* Enable multi-gpu training in find_learning_rate.py

* Added test for multi-gpu training

* Minor changes in tests

* Change in indentation

* Changes in testing multi-gpu usage

Removed new class for testing multi-gpu usage and moved the testing function to existing class `TestFindLearningRate`

* Changes in find_learning_rate_test.py

* Minor changes in find_learning_rate_test.py

* Remove redundant code

Author: gurunathparasaram

allennlp/commands/find_learning_rate.py

@@ -170,7 +170,12 @@ def find_learning_rate_model(params: Params, serialization_dir: str,
 
     prepare_environment(params)
 
-    check_for_gpu(params.get('trainer').get('cuda_device', -1))
+    cuda_device = params.params.get('trainer').get('cuda_device', -1)
+    if isinstance(cuda_device, list):
+        for device in cuda_device:
+            check_for_gpu(device)
+    else:
+        check_for_gpu(cuda_device)
 
     all_datasets = datasets_from_params(params)
     datasets_for_vocab_creation = set(params.pop("datasets_for_vocab_creation", all_datasets))

allennlp/tests/commands/find_learning_rate_test.py

@@ -3,6 +3,8 @@
 import os
 import pytest
 
+import torch
+
 from allennlp.common import Params
 from allennlp.data import Vocabulary, DataIterator
 from allennlp.models import Model
@@ -12,6 +14,7 @@
 from allennlp.commands.find_learning_rate import search_learning_rate, \
     find_learning_rate_from_args, find_learning_rate_model, FindLearningRate
 
+
 class TestFindLearningRate(AllenNlpTestCase):
 
     def setUp(self):
@@ -44,7 +47,8 @@ def setUp(self):
                     })
 
     def test_find_learning_rate(self):
-        find_learning_rate_model(self.params(), os.path.join(self.TEST_DIR, 'test_find_learning_rate'),
+        find_learning_rate_model(self.params(),
+                                 os.path.join(self.TEST_DIR, 'test_find_learning_rate'),
                                  start_lr=1e-5,
                                  end_lr=1,
                                  num_batches=100,
@@ -89,7 +93,6 @@ def test_find_learning_rate(self):
                                  stopping_factor=None,
                                  force=True)
 
-
     def test_find_learning_rate_args(self):
         parser = argparse.ArgumentParser(description="Testing")
         subparsers = parser.add_subparsers(title='Commands', metavar='')
@@ -115,6 +118,21 @@ def test_find_learning_rate_args(self):
             assert cm.exception.code == 2  # argparse code for incorrect usage
 
 
+    @pytest.mark.skipif(torch.cuda.device_count() < 2,
+                        reason="Need multiple GPUs.")
+    def test_find_learning_rate_multi_gpu(self):
+        params = self.params()
+        params["trainer"]["cuda_device"] = [0, 1]
+        find_learning_rate_model(params,
+                                 os.path.join(self.TEST_DIR, 'test_find_learning_rate_multi_gpu'),
+                                 start_lr=1e-5,
+                                 end_lr=1,
+                                 num_batches=100,
+                                 linear_steps=True,
+                                 stopping_factor=None,
+                                 force=False)
+
+
 class TestSearchLearningRate(AllenNlpTestCase):
 
     def setUp(self):
@@ -144,7 +162,7 @@ def setUp(self):
                     "num_epochs": 2,
                     "optimizer": "adam"
                 }
-        })
+            })
         all_datasets = datasets_from_params(params)
         vocab = Vocabulary.from_params(
             params.pop("vocabulary", {}),
@@ -159,12 +177,12 @@ def setUp(self):
         serialization_dir = os.path.join(self.TEST_DIR, 'test_search_learning_rate')
 
         self.trainer = Trainer.from_params(model,
-                                      serialization_dir,
-                                      iterator,
-                                      train_data,
-                                      params=trainer_params,
-                                      validation_data=None,
-                                      validation_iterator=None)
+                                           serialization_dir,
+                                           iterator,
+                                           train_data,
+                                           params=trainer_params,
+                                           validation_data=None,
+                                           validation_iterator=None)
 
     def test_search_learning_rate_with_num_batches_less_than_ten(self):
         with pytest.raises(ConfigurationError):
@@ -175,6 +193,7 @@ def test_search_learning_rate_linear_steps(self):
         assert len(learning_rates_losses) > 1
 
     def test_search_learning_rate_without_stopping_factor(self):
-        learning_rates, losses = search_learning_rate(self.trainer, num_batches=100, stopping_factor=None)
+        learning_rates, losses = search_learning_rate(self.trainer, num_batches=100,
+                                                      stopping_factor=None)
         assert len(learning_rates) == 101
         assert len(losses) == 101