Caffe2 - C++ API
A deep learning, cross platform ML framework
inspect_gpus.cc
1 
17 #include <cuda_runtime.h>
18 
19 #include <sstream>
20 #include <vector>
21 
22 #include "caffe2/core/common_gpu.h"
23 #include "caffe2/core/init.h"
24 #include "caffe2/core/logging.h"
25 
26 using std::vector;
27 
28 CAFFE2_DECLARE_int(caffe2_log_level);
29 
30 int main(int argc, char** argv) {
31  caffe2::GlobalInit(&argc, &argv);
33  "Inspects the GPUs on the current machine and prints out their details "
34  "provided by cuda.");
35 
36  int gpu_count;
37  CUDA_ENFORCE(cudaGetDeviceCount(&gpu_count));
38  for (int i = 0; i < gpu_count; ++i) {
39  LOG(INFO) << "Querying device ID = " << i;
41  }
42 
43  vector<vector<bool> > access_pattern;
44  CAFFE_ENFORCE(caffe2::GetCudaPeerAccessPattern(&access_pattern));
45 
46  std::stringstream sstream;
47  // Find topology
48  for (int i = 0; i < gpu_count; ++i) {
49  for (int j = 0; j < gpu_count; ++j) {
50  sstream << (access_pattern[i][j] ? "+" : "-") << " ";
51  }
52  sstream << std::endl;
53  }
54  LOG(INFO) << "Access pattern: " << std::endl << sstream.str();
55 
56  return 0;
57 }
bool GlobalInit(int *pargc, char ***pargv)
Initialize the global environment of caffe2.
Definition: init.cc:34
void DeviceQuery(const int device)
Runs a device query function and prints out the results to LOG(INFO).
Definition: common_gpu.cc:199
bool GetCudaPeerAccessPattern(vector< vector< bool > > *pattern)
Return a peer access pattern by returning a matrix (in the format of a nested vector) of boolean valu...
Definition: common_gpu.cc:234
void SetUsageMessage(const string &str)
Sets the usage message when a commandline tool is called with "--help".
Definition: flags.cc:68