libtorch c++调用 （七）选择显卡运行torchscript

1、选择显卡运行

#include 

std::string filename = "centernet.pt"//模型路径
int gpu_id = 1;									//gpu　id 0代表第一块可见gpu
cudaSetDevice(gpu_id);					//切换显卡
torch::jit::script::Module module = torch::jit::load(filename,torch::Device(torch::DeviceType::CUDA,gpu_id));//加载模型

libtorch 加载torchscript模型有三个重载函数

/// Loads a serialized `Module` from the given `istream`.
///
/// The istream must contain a serialized `Module`, exported via
/// `torch::jit::ExportModule` in C++.
TORCH_API Module load(
    std::istream& in,
    c10::optional device = c10::nullopt,
    ExtraFilesMap& extra_files = default_extra_files);

/// Loads a serialized `Module` from the given `filename`.
///
/// The file stored at the location given in `filename` must contain a
/// serialized `Module`, exported either via `ScriptModule.save()` in
/// Python or `torch::jit::ExportModule` in C++.
TORCH_API Module load(
    const std::string& filename,
    c10::optional device = c10::nullopt,
    ExtraFilesMap& extra_files = default_extra_files);

/// Loads a serialized `Module` from the given `rai`.
///
/// The reader adapter, which is for customized input stream, must contain a
/// serialized `Module`, exported either via `ScriptModule.save()` in
/// Python or `torch::jit::ExportModule` in C++.
TORCH_API Module load(
    std::unique_ptr rai,
    c10::optional device = c10::nullopt,
    ExtraFilesMap& extra_files = default_extra_files);

目前我是从文件加载模型，用第二个函数，选择设备这里主要关注第二个参数

c10::optional device

这里我们需要构造一个device类传入，我们看Device类定义

Device(DeviceType type, DeviceIndex index = -1)

这里很显然第一个是设备类型，第二个是设备索引
第一个是枚举类：我们选择torch::DeviceType::CUDA　　也就是nvidia显卡计算平台
第一个就是显卡id 　我们填0代表第一块显卡,１代表第二块显卡

 

参考：libtorch选择显卡运行torchscript

 

2、选择指定显卡运行

方法

以指定使用显卡1为例：

//首先包含头文件
#include 

//
// 程序开始植入代码
int8_t cuda_index = 1;
c10::cuda::CUDAGuard device_guard(cuda_index); //指定默认的显卡索引
.... 
//将模型加载至显卡中
module.to(at::kCUDA);

//将Tensor数据加载至显卡，
img_tensor = img_tensor.to(at::kCUDA);

参考：libtorch（pytorch C++） - 使用指定的显卡