[API2.0] add Device api (set_device and get_device)(#26103)

paddle/fluid/pybind/pybind.cc

@@ -1320,6 +1320,8 @@ All parameter, weight, gradient are variables in Paddle.
 #endif
            })
 #ifdef PADDLE_WITH_CUDA
+      .def("get_device_id",
+           [](const platform::CUDAPlace &self) { return self.GetDeviceId(); })
       .def("_type", &PlaceIndex<platform::CUDAPlace>)
       .def("_equals", &IsSamePlace<platform::CUDAPlace, platform::Place>)
       .def("_equals", &IsSamePlace<platform::CUDAPlace, platform::CUDAPlace>)

python/paddle/__init__.py

@@ -42,6 +42,7 @@ import paddle.nn
 import paddle.distributed.fleet
 import paddle.optimizer
 import paddle.metric
+import paddle.device
 import paddle.incubate.complex as complex
 
 # TODO: define alias in tensor and framework directory
@@ -231,6 +232,9 @@ from .tensor.stat import reduce_mean  #DEFINE_ALIAS
 from .tensor.stat import std  #DEFINE_ALIAS
 from .tensor.stat import var  #DEFINE_ALIAS
 from .fluid.data import data
+from .device import set_device
+from .device import get_device
+# from .tensor.tensor import Tensor        #DEFINE_ALIAS
 # from .tensor.tensor import LoDTensor        #DEFINE_ALIAS
 # from .tensor.tensor import LoDTensorArray        #DEFINE_ALIAS
 

python/paddle/device.py

@@ -12,11 +12,83 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# TODO: define the functions to manipulate devices 
-# __all__ = ['cpu_places',
-#            'CPUPlace',
-#            'cuda_pinned_places',
-#            'cuda_places',
-#            'CUDAPinnedPlace',
-#            'CUDAPlace',
-#            'is_compiled_with_cuda']
+from paddle.fluid import core
+from paddle.fluid import framework
+import re
+__all__ = [
+    'set_device',
+    'get_device'
+    #            'cpu_places',
+    #            'CPUPlace',
+    #            'cuda_pinned_places',
+    #            'cuda_places',
+    #            'CUDAPinnedPlace',
+    #            'CUDAPlace',
+    #            'is_compiled_with_cuda'
+]
+
+
+def set_device(device):
+    """
+    Paddle supports running calculations on various types of devices, including CPU and GPU.
+    They are represented by string identifiers. This function can specify the global device
+    which the OP will run.
+
+    Parameters:
+        device(str): This parameter determines the specific running device.
+            It can be ``cpu`` or ``gpu:0``. When ``device`` is ``cpu``, the
+            program is running on the cpu. When ``device`` is ``gpu``, the
+            program is running ont the gpu.
+    Examples:
+
+     .. code-block:: python
+            
+        import paddle
+        paddle.enable_imperative()
+        paddle.fluid.dygraph.set_device("gpu:0")
+        x1 = paddle.ones(name='x1', shape=[1, 2], dtype='int32')
+        x2 = paddle.zeros(name='x2', shape=[1, 2], dtype='int32')
+        data = paddle.stack([x1,x2], axis=1)
+    """
+    lower_device = device.lower()
+    if lower_device == 'cpu':
+        place = core.CPUPlace()
+        framework._set_expected_place(place)
+    else:
+        avaliable_device = ((lower_device == 'cpu') or
+                            re.match(r'gpu:\d+', lower_device))
+        if not avaliable_device:
+            raise ValueError(
+                "The device must be a string which is like 'cpu' or 'gpu:0'")
+        device_info_list = device.split(':', 1)
+        device_id = device_info_list[1]
+        device_id = int(device_id)
+        place = core.CUDAPlace(device_id)
+        framework._set_expected_place(place)
+
+
+def get_device():
+    """
+    This funciton can get the current global device of the program is running.
+    It's a string which is like 'cpu' and 'gpu:0'. if the global device is not
+    set, it will return a string which is 'gpu:0' when cuda is avaliable or it 
+    will return a string which is 'cpu' when cuda is not avaliable.
+
+    Examples:
+
+     .. code-block:: python
+            
+        import paddle
+        paddle.enable_imperative()
+        device = paddle.fluid.dygraph.get_device()
+
+    """
+    device = ''
+    place = framework._current_expected_place()
+    if isinstance(place, core.CPUPlace):
+        device = 'cpu'
+    elif isinstance(place, core.CUDAPlace):
+        device_id = place.get_device_id()
+        device = 'gpu:' + str(device_id)
+
+    return device

python/paddle/fluid/dygraph/base.py

@@ -26,13 +26,8 @@ import objgraph
 from ..data_feeder import convert_dtype
 
 __all__ = [
-    'no_grad',
-    'grad',
-    'guard',
-    'enable_dygraph',
-    'disable_dygraph',
-    'enabled',
-    'to_variable',
+    'no_grad', 'grad', 'guard', 'enable_dygraph', 'disable_dygraph', 'enabled',
+    'to_variable'
 ]
 
 
@@ -285,12 +280,11 @@ def guard(place=None):
     tracer = Tracer()
     VarBase = core.VarBase
 
-    if place is None:
-        if core.is_compiled_with_cuda():
-            place = core.CUDAPlace(0)
+    if place is not None:
+        expected_place = place
     else:
-            place = core.CPUPlace()
-    tracer._expected_place = place
+        expected_place = framework._current_expected_place()
+    tracer._expected_place = expected_place
 
     with framework.program_guard(train, startup):
         with framework.unique_name.guard():

python/paddle/fluid/executor.py

@@ -31,6 +31,7 @@ from .. import compat as cpt
 from .trainer_factory import TrainerFactory
 from .trainer_factory import FetchHandlerMonitor
 import copy
+from . import framework
 from .incubate.checkpoint import auto_checkpoint as acp
 
 __all__ = ['Executor', 'global_scope', 'scope_guard']
@@ -544,10 +545,8 @@ class Executor(object):
 
     def __init__(self, place=None):
         if place is None:
-            if core.is_compiled_with_cuda():
-                self.place = core.CUDAPlace(0)
-            else:
-                self.place = core.CPUPlace()
+            expected_place = framework._current_expected_place()
+            self.place = expected_place
         else:
             self.place = place
         self.program_caches = dict()

python/paddle/fluid/framework.py

@@ -64,7 +64,7 @@ ZERO_VAR_SUFFIX = core.kZeroVarSuffix()
 CONTROL_DEP_VAR_PREFIX = core.kControlDepVarName()
 
 _dygraph_tracer_ = None
-_dygraph_current_expected_place_ = None
+_global_expected_place_ = None
 _current_device = None
 global_prog_seed = 0
 
@@ -247,7 +247,26 @@ def _dygraph_tracer():
 
 
 def _current_expected_place():
-    return _dygraph_current_expected_place_
+    global _global_expected_place_
+    if _global_expected_place_ is None:
+        if core.is_compiled_with_cuda():
+            _global_expected_place_ = core.CUDAPlace(0)
+        else:
+            _global_expected_place_ = core.CPUPlace()
+
+    return _global_expected_place_
+
+
+def _set_dygraph_tracer_expected_place(place):
+    global _dygraph_tracer_
+    if _dygraph_tracer_ is not None:
+        _dygraph_tracer_._expected_place = place
+
+
+def _set_expected_place(place):
+    global _global_expected_place_
+    _global_expected_place_ = place
+    _set_dygraph_tracer_expected_place(place)
 
 
 # TODO(zhiqiu): remove this function.
@@ -5417,14 +5436,14 @@ def _dygraph_guard(tracer):
 
 @signature_safe_contextmanager
 def _dygraph_place_guard(place):
-    global _dygraph_current_expected_place_
-    tmp_place = _dygraph_current_expected_place_
-    _dygraph_current_expected_place_ = place
+    global _global_expected_place_
+    tmp_place = _global_expected_place_
+    _global_expected_place_ = place
 
     try:
         yield
     finally:
-        _dygraph_current_expected_place_ = tmp_place
+        _global_expected_place_ = tmp_place
 
 
 def load_op_library(lib_filename):

python/paddle/fluid/tests/unittests/test_device.py

+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import unittest
+from op_test import OpTest
+
+import numpy as np
+import paddle.fluid as fluid
+import paddle.fluid.core as core
+import paddle.fluid.framework as framework
+import warnings
+import paddle
+
+
+class TestStaticDeviceManage(unittest.TestCase):
+    def test_cpu_device(self):
+        paddle.set_device('cpu')
+        out1 = paddle.zeros(shape=[1, 3], dtype='float32')
+        out2 = paddle.ones(shape=[1, 3], dtype='float32')
+        out3 = paddle.concat(x=[out1, out2], axis=0)
+        exe = paddle.fluid.Executor()
+        exe.run(paddle.fluid.default_startup_program())
+        res = exe.run(fetch_list=[out3])
+        device = paddle.get_device()
+        self.assertEqual(isinstance(exe.place, core.CPUPlace), True)
+        self.assertEqual(device, "cpu")
+
+    def test_gpu_device(self):
+        if core.is_compiled_with_cuda():
+            out1 = paddle.zeros(shape=[1, 3], dtype='float32')
+            out2 = paddle.ones(shape=[1, 3], dtype='float32')
+            out3 = paddle.concat(x=[out1, out2], axis=0)
+            paddle.set_device('gpu:0')
+            exe = paddle.fluid.Executor()
+            exe.run(paddle.fluid.default_startup_program())
+            res = exe.run(fetch_list=[out3])
+            device = paddle.get_device()
+            self.assertEqual(isinstance(exe.place, core.CUDAPlace), True)
+            self.assertEqual(device, "gpu:0")
+
+
+class TestImperativeDeviceManage(unittest.TestCase):
+    def test_cpu(self):
+        with fluid.dygraph.guard():
+            paddle.set_device('cpu')
+            out1 = paddle.zeros(shape=[1, 3], dtype='float32')
+            out2 = paddle.ones(shape=[1, 3], dtype='float32')
+            out3 = paddle.concat(x=[out1, out2], axis=0)
+            device = paddle.get_device()
+            self.assertEqual(
+                isinstance(framework._current_expected_place(), core.CPUPlace),
+                True)
+            self.assertEqual(device, "cpu")
+
+    def test_gpu(self):
+        if core.is_compiled_with_cuda():
+            with fluid.dygraph.guard():
+                paddle.set_device('gpu:0')
+                out1 = paddle.zeros(shape=[1, 3], dtype='float32')
+                out2 = paddle.ones(shape=[1, 3], dtype='float32')
+                out3 = paddle.concat(x=[out1, out2], axis=0)
+                device = paddle.get_device()
+                self.assertEqual(
+                    isinstance(framework._current_expected_place(),
+                               core.CUDAPlace), True)
+                self.assertEqual(device, "gpu:0")
+
+
+if __name__ == '__main__':
+    unittest.main()