Repository: keras-team/autokeras

Issue:
Parallel on Multiple GPUs
I tried fitting the data on amazon p2.8xlarge instance(which has 8 GPUs), but it just fills the VRAM of one GPU and crashes with the cuda out of memory error. Is there any way of using all 8 GPUs?


Discussion:
+1.

Only running on one GPU despite having four GPU on my server.


More general question:
Can we run autokeras models (e.g., ImageClassifier) on a cluster of distributed machines with each having multiple GPUs? 
I tried to run it in this fashion, but, was not successful. It seems like it only runs on one GPU in a single cluster node. 

More clarification on this feature would be appreciated a lot.
Thanks a lot 
Strongly looking forward to this feature!

Resolving patch:
diff --git a/autokeras/nn/graph.py b/autokeras/nn/graph.py
--- a/autokeras/nn/graph.py
+++ b/autokeras/nn/graph.py
@@ -675,14 +675,13 @@ def forward(self, input_tensor):
         for v in topo_node_list:
             for u, layer_id in self.graph.reverse_adj_list[v]:
                 layer = self.graph.layer_list[layer_id]
-                torch_layer = self.layers[layer_id]
+                torch_layer = list(self.modules())[layer_id + 1]
 
                 if isinstance(layer, (StubAdd, StubConcatenate)):
                     edge_input_tensor = list(map(lambda x: node_list[x],
                                                  self.graph.layer_id_to_input_node_ids[layer_id]))
                 else:
                     edge_input_tensor = node_list[u]
-
                 temp_tensor = torch_layer(edge_input_tensor)
                 node_list[v] = temp_tensor
         return node_list[output_id]
diff --git a/autokeras/nn/model_trainer.py b/autokeras/nn/model_trainer.py
--- a/autokeras/nn/model_trainer.py
+++ b/autokeras/nn/model_trainer.py
@@ -39,7 +39,7 @@ def __init__(self,
                  metric=None,
                  verbose=False,
                  device=None):
-        if device is not None:
+        if device:
             self.device = device
         else:
             self.device = get_device()
@@ -84,6 +84,8 @@ class ModelTrainer(ModelTrainerBase):
     def __init__(self, model, path, **kwargs):
         super().__init__(**kwargs)
         self.model = model
+        if torch.cuda.device_count() > 1:
+            self.model = torch.nn.DataParallel(self.model)
         self.model.to(self.device)
         self.optimizer = None
         self.early_stop = None
diff --git a/autokeras/utils.py b/autokeras/utils.py
--- a/autokeras/utils.py
+++ b/autokeras/utils.py
@@ -52,39 +52,48 @@ def pickle_to_file(obj, path):
     pickle.dump(obj, open(path, 'wb'))
 
 
-# TODO cannot detect nvidia-smi in Windows normally. We need a fall back for windows
 def get_device():
     """ If CUDA is available, use CUDA device, else use CPU device.
-    When choosing from CUDA devices, this function will choose the one with max memory available.
-    Returns: string device name.
+    Returns: string device name
     """
-    # TODO: could use gputil in the future
-    device = 'cpu'
-    if torch.cuda.is_available():
-        try:
-            # smi_out=
-            #       Free                 : xxxxxx MiB
-            #       Free                 : xxxxxx MiB
-            #                      ....
-            smi_out = subprocess.check_output('nvidia-smi -q -d Memory | grep -A4 GPU|grep Free', shell=True)
-            if isinstance(smi_out, bytes):
-                smi_out = smi_out.decode('utf-8')
-        except subprocess.SubprocessError:
-            warnings.warn('Cuda device successfully detected. However, nvidia-smi cannot be invoked')
-            return 'cpu'
-        visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', '').split(',')
-        if len(visible_devices) == 1 and visible_devices[0] == '':
-            visible_devices = []
-        visible_devices = [int(x) for x in visible_devices]
-        memory_available = [int(x.split()[2]) for x in smi_out.splitlines()]
-        for cuda_index, _ in enumerate(memory_available):
-            if cuda_index not in visible_devices and visible_devices:
-                memory_available[cuda_index] = 0
-
-        if memory_available:
-            if max(memory_available) != 0:
-                device = 'cuda:' + str(memory_available.index(max(memory_available)))
-    return device
+    return 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+# # TODO cannot detect nvidia-smi in Windows normally. We need a fall back for windows
+# def get_device():
+#     """ If CUDA is available, use CUDA device, else use CPU device.
+#     When choosing from CUDA devices, this function will choose the one with max memory available.
+#     Returns: string device name.
+#     """
+#     # TODO: could use gputil in the future
+#     device = 'cpu'
+#     if torch.cuda.is_available():
+#         try:
+#             # smi_out=
+#             #       Free                 : xxxxxx MiB
+#             #       Free                 : xxxxxx MiB
+#             #                      ....
+#             smi_out = subprocess.check_output('nvidia-smi -q -d Memory | grep -A4 GPU|grep Free', shell=True)
+#             if isinstance(smi_out, bytes):
+#                 smi_out = smi_out.decode('utf-8')
+#         except subprocess.SubprocessError:
+#             warnings.warn('Cuda device successfully detected. However, nvidia-smi cannot be invoked')
+#             return 'cpu'
+#         visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', '').split(',')
+#         if len(visible_devices) == 1 and visible_devices[0] == '':
+#             visible_devices = []
+#         visible_devices = [int(x) for x in visible_devices]
+#         memory_available = [int(x.split()[2]) for x in smi_out.splitlines()]
+#         for cuda_index, _ in enumerate(memory_available):
+#             if cuda_index not in visible_devices and visible_devices:
+#                 memory_available[cuda_index] = 0
+#         memory_available = list(filter(lambda a: a != 2, memory_available))
+#         if memory_available:
+#             if max(memory_available) != 0 and torch.cuda.device_count() == 1:
+#                 device = 'cuda:' + str(memory_available.index(max(memory_available)))
+#             elif max(memory_available) != 0 and torch.cuda.device_count() > 1:
+#                 device = 'cuda:0'
+#     return device
 
 
 def temp_path_generator():


Test patch:
diff --git a/tests/common.py b/tests/common.py
--- a/tests/common.py
+++ b/tests/common.py
@@ -316,9 +316,9 @@ def mock_out_of_memory_train(**kwargs):
     raise RuntimeError('CUDA: out of memory.')
 
 
-def mock_nvidia_smi_output(*arg, **kwargs):
-    return \
-        '    Free                        : 1 MiB \n' \
-        '    Free                        : 11176 MiB \n' \
-        '    Free                        : 1 MiB \n' \
-        '    Free                        : 1 MiB'
+# def mock_nvidia_smi_output(*arg, **kwargs):
+#     return \
+#         '    Free                        : 1 MiB \n' \
+#         '    Free                        : 11176 MiB \n' \
+#         '    Free                        : 1 MiB \n' \
+#         '    Free                        : 1 MiB'
diff --git a/tests/test_utils.py b/tests/test_utils.py
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,12 +1,11 @@
-import numpy as np
-import subprocess
 from unittest.mock import patch
 
+import numpy as np
+
 from autokeras.constant import Constant
-from autokeras.utils import rand_temp_folder_generator, download_file, get_system, get_device, \
-    compute_image_resize_params, \
+from autokeras.utils import rand_temp_folder_generator, download_file, get_system, compute_image_resize_params, \
     resize_image_data, temp_path_generator
-from tests.common import clean_dir, TEST_TEMP_DIR, mock_nvidia_smi_output, TEST_TEMP_KERAS_DIR
+from tests.common import clean_dir, TEST_TEMP_DIR, TEST_TEMP_KERAS_DIR
 
 
 # This method will be used by the mock to replace requests.get
@@ -90,19 +89,19 @@ def test_get_system():
         sys_name == Constant.SYS_WINDOWS
 
 
-@patch('torch.cuda.is_available')
-@patch('subprocess.check_output')
-def test_get_device(mock_check_output, mock_is_available):
-    mock_check_output.return_value = mock_nvidia_smi_output()
-    mock_is_available.return_value = True
-    device = get_device()
-    assert device == 'cuda:1'
-
-    mock_check_output.side_effect = subprocess.SubprocessError
-    device = get_device()
-    assert device == 'cpu'
-
-    mock_is_available.return_value = False
-    mock_check_output.return_value = ''
-    device = get_device()
-    assert device == 'cpu'
+# @patch('torch.cuda.is_available')
+# @patch('subprocess.check_output')
+# def test_get_device(mock_check_output, mock_is_available):
+#     mock_check_output.return_value = mock_nvidia_smi_output()
+#     mock_is_available.return_value = True
+#     device = get_device()
+#     assert device == 'cuda:1'
+#
+#     mock_check_output.side_effect = subprocess.SubprocessError
+#     device = get_device()
+#     assert device == 'cpu'
+#
+#     mock_is_available.return_value = False
+#     mock_check_output.return_value = ''
+#     device = get_device()
+#     assert device == 'cpu'