diff options
| author | Evan Shelhamer <shelhamer@imaginarynumber.net> | 2014-04-26 15:22:42 -0700 |
|---|---|---|
| committer | Evan Shelhamer <shelhamer@imaginarynumber.net> | 2014-04-26 15:22:42 -0700 |
| commit | 0b885b9ba4f784fa3aad6444eb08cf592e67017b (patch) | |
| tree | 21463a11a81e060510c8411ff228579fefeb6d79 /tools/test_net.cpp | |
| parent | a154ebcab3b9ded8b2bfdf5708ead1fd645644f5 (diff) | |
| download | caffeonacl-0b885b9ba4f784fa3aad6444eb08cf592e67017b.tar.gz caffeonacl-0b885b9ba4f784fa3aad6444eb08cf592e67017b.tar.bz2 caffeonacl-0b885b9ba4f784fa3aad6444eb08cf592e67017b.zip | |
add device id arg to test_net (fix #232)
Diffstat (limited to 'tools/test_net.cpp')
| -rw-r--r-- | tools/test_net.cpp | 11 |
1 files changed, 7 insertions, 4 deletions
diff --git a/tools/test_net.cpp b/tools/test_net.cpp index 559fa730..69c774a6 100644 --- a/tools/test_net.cpp +++ b/tools/test_net.cpp @@ -17,18 +17,21 @@ using namespace caffe; // NOLINT(build/namespaces) int main(int argc, char** argv) { - if (argc < 4 || argc > 5) { + if (argc < 4 || argc > 6) { LOG(ERROR) << "test_net net_proto pretrained_net_proto iterations " - << "[CPU/GPU]"; + << "[CPU/GPU] [Device ID]"; return 1; } - cudaSetDevice(0); Caffe::set_phase(Caffe::TEST); - if (argc == 5 && strcmp(argv[4], "GPU") == 0) { + if (argc >= 5 && strcmp(argv[4], "GPU") == 0) { LOG(ERROR) << "Using GPU"; Caffe::set_mode(Caffe::GPU); + if (argc == 6) { + int device_id = atoi(argv[5]); + Caffe::SetDevice(device_id); + } } else { LOG(ERROR) << "Using CPU"; Caffe::set_mode(Caffe::CPU); |