"fix gpu init"

paddle/framework/init.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 #include <string.h>  // for strdup
 #include <algorithm>
+#include <stdexcept>
 #include <string>
 
 #include "paddle/framework/init.h"
@@ -46,17 +47,23 @@ void InitDevices() {
 
   std::vector<platform::Place> places;
   places.emplace_back(platform::CPUPlace());
+  int count = 0;
 
 #ifdef PADDLE_WITH_CUDA
-  int count = platform::GetCUDADeviceCount();
-  for (int i = 0; i < count; ++i) {
-    places.emplace_back(platform::CUDAPlace(i));
+  try {
+    count = platform::GetCUDADeviceCount();
+  } catch (const std::exception &exp) {
+    LOG(WARNING) << "Compiled with WITH_GPU, but no GPU found in runtime.";
   }
 #else
   LOG(WARNING)
-      << "'GPU' is not supported, Please re-compile with WITH_GPU option";
+      << "'CUDA' is not supported, Please re-compile with WITH_GPU option";
 #endif
 
+  for (int i = 0; i < count; ++i) {
+    places.emplace_back(platform::CUDAPlace(i));
+  }
+
   platform::DeviceContextPool::Init(places);
 }
 

paddle/framework/init_test.cc

@@ -20,7 +20,26 @@ TEST(InitDevices, CPU) {
   using paddle::framework::InitDevices;
   using paddle::platform::DeviceContextPool;
 
+#ifndef PADDLE_WITH_CUDA
   InitDevices();
   DeviceContextPool& pool = DeviceContextPool::Instance();
-  ASSERT_GE(pool.size(), 1U);
+  ASSERT_EQ(pool.size(), 1U);
+#endif
+}
+
+TEST(InitDevices, CUDA) {
+  using paddle::framework::InitDevices;
+  using paddle::platform::DeviceContextPool;
+
+  int count = 0;
+  try {
+    count = paddle::platform::GetCUDADeviceCount();
+  } catch (const std::exception& exp) {
+  }
+
+#ifdef PADDLE_WITH_CUDA
+  InitDevices();
+  DeviceContextPool& pool = DeviceContextPool::Instance();
+  ASSERT_EQ(pool.size(), 1U + static_cast<unsigned>(count));
+#endif
 }

python/paddle/v2/fluid/__init__.py

@@ -91,9 +91,23 @@ def __bootstrap__():
     ]
     if core.is_compile_gpu():
         read_env_flags += ['fraction_of_gpu_memory_to_use', 'op_sync']
+
     core.init_gflags([sys.argv[0]] +
                      ["--tryfromenv=" + ",".join(read_env_flags)])
     core.init_glog(sys.argv[0])
+
+    gpu_devices = os.getenv("CUDA_VISIBLE_DEVICES", '')
+    if core.is_compile_gpu():
+        if len(gpu_devices.split(",")) >= 1:
+            print(
+                'WARNING: CUDA_VISIBLE_DEVICES set to {0}, not empty . The computation '
+                'speed will not be optimized if you use multi-gpu. It will '
+                'fail if this PaddlePaddle binary is compiled without GPU option'
+                .format(gpu_devices),
+                file=sys.stderr)
+        else:
+            gpu_devices = "0"
+    os.environ['CUDA_VISIBLE_DEVICES'] = gpu_devices
     core.init_devices()