blob: c8caf5ac58ef4f5cd949ec60ff48041fede26284 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
|
// The main caffe test code. Your test cpp code should include this hpp
// to allow a main function to be compiled into the binary.
#include "caffe/caffe.hpp"
#include "caffe/test/test_caffe_main.hpp"
namespace caffe {
#ifndef CPU_ONLY
cudaDeviceProp CAFFE_TEST_CUDA_PROP;
#endif
}
#ifndef CPU_ONLY
using caffe::CAFFE_TEST_CUDA_PROP;
#endif
int main(int argc, char** argv) {
::testing::InitGoogleTest(&argc, argv);
caffe::GlobalInit(&argc, &argv);
#ifndef CPU_ONLY
// Before starting testing, let's first print out a few cuda defice info.
int device;
cudaGetDeviceCount(&device);
cout << "Cuda number of devices: " << device << endl;
if (argc > 1) {
// Use the given device
device = atoi(argv[1]);
cudaSetDevice(device);
cout << "Setting to use device " << device << endl;
} else if (CUDA_TEST_DEVICE >= 0) {
// Use the device assigned in build configuration; but with a lower priority
device = CUDA_TEST_DEVICE;
}
cudaGetDevice(&device);
cout << "Current device id: " << device << endl;
cudaGetDeviceProperties(&CAFFE_TEST_CUDA_PROP, device);
#endif
// invoke the test.
return RUN_ALL_TESTS();
}
|