How to specify CUDA device number in C++?

C++
1

Hello,
I have four different libtorch services and also four GPU’s available. In PyTorch its is quite easy to state the GPU device, how do I do it in C++ so that each service is assigned to a different GPU?

Thank,

2

The .to() operation should also work in C++ as described here.

3

Thanks, for reference:
Change:

torch::jit::script::Module module = torch::jit::load(s_model_name, torch::kCUDA);

To:

torch::jit::script::Module module = torch::jit::load(s_model_name, torch::kCUDA);
const std::string device_string = "cuda:2";
module.to(device_string);

And:

input_tensor = input_tensor.to(at::kCUDA);

To:

input_tensor = input_tensor.to(device_string);

image
image1185×400 36.1 KB

3 Likes
4

nice! very simple, thank you very much!

5

Confirm that this works!

6

Maybe too late, but after reading the correction, I suggest to not use raw string to specify your device number, but instead use a little helper like:

auto get_cuda_device(const int number) -> std::string
{
     assert(number >= 0 && number <= torch::cuda::device_count());
     return "cuda:" + std::to_string(number);
}

just replace:

const std::string device_string = "cuda:2";

by

auto device_string = get_cuda_device(2):

This way of doing the code prevents:

  1. Problems with hardware changes
  2. Typo problems when using several gpu numbers.