src/caffe/test/test_caffe_main.hpp

   1 // Copyright 2013 Yangqing Jia
   2
   3 // The main caffe test code. Your test cpp code should include this hpp
   4 // to allow a main function to be compiled into the binary.
   5 #ifndef CAFFE_TEST_TEST_CAFFE_MAIN_HPP_
   6 #define CAFFE_TEST_TEST_CAFFE_MAIN_HPP_
   7
   8 #include <cuda_runtime.h>
   9 #include <glog/logging.h>
  10 #include <gtest/gtest.h>
  11
  12 #include <cstdlib>
  13 #include <cstdio>
  14 #include <iostream>
  15
  16
  17 namespace caffe {
  18
  19 cudaDeviceProp CAFFE_TEST_CUDA_PROP;
  20
  21 }  // namespace caffe
  22
  23 using namespace caffe;
  24 using namespace std;
  25
  26 int main(int argc, char** argv) {
  27   ::testing::InitGoogleTest(&argc, argv);
  28   ::google::InitGoogleLogging(argv[0]);
  29   // Before starting testing, let's first print out a few cuda defice info.
  30   int device;
  31   cudaGetDeviceCount(&device);
  32   cout << "Cuda number of devices: " << device << endl;
  33   if (argc > 1) {
  34     // Use the given device
  35     device = atoi(argv[1]);
  36     cudaSetDevice(device);
  37     cout << "Setting to use device " << device << endl;
  38   }
  39   cudaGetDevice(&device);
  40   cout << "Current device id: " << device << endl;
  41   cudaGetDeviceProperties(&CAFFE_TEST_CUDA_PROP, device);
  42   // invoke the test.
  43   return RUN_ALL_TESTS();
  44 }
  45
  46 #endif  // CAFFE_TEST_TEST_CAFFE_MAIN_HPP_