windows libtorch cuda sample
时间: 2025-01-09 07:16:02 浏览: 6
### 使用LibTorch与CUDA在Windows上的示例代码
为了展示如何在Windows环境下利用LibTorch和CUDA编写程序,下面提供了一个简单的例子来说明设置环境以及实现基本操作的方法。
#### 安装依赖项
确保已经安装了适用于Windows子系统Linux(WSL)的NVIDIA CUDA工具包[^1]。这一步骤对于支持GPU加速至关重要。接着按照官方指南完成PyTorch C++前端(LibTorch)的下载并配置开发环境。
#### 创建CMakeLists.txt文件
创建项目结构时,在根目录下放置`CMakeLists.txt`用于构建过程管理:
```cmake
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(example)
find_package(Torch REQUIRED)
add_executable(${PROJECT_NAME} main.cpp)
target_link_libraries(${PROJECT_NAME} "${TORCH_LIBRARIES}")
set_property(TARGET ${PROJECT_NAME} PROPERTY CXX_STANDARD 14)
if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DVERSION_INFO=\\\"${PROJECT_VERSION}\\\"")
endif()
```
#### 编写main.cpp源码
接下来定义应用程序逻辑入口点——`main()`函数,并尝试加载预训练模型到GPU设备上执行推理任务:
```cpp
#include <torch/torch.h>
#include <iostream>
int main() {
// 检查是否有可用的CUDA设备
if (!torch::cuda::is_available()) {
std::cerr << "CUDA not available." << std::endl;
return -1;
}
// 加载ResNet-18预训练模型并将其实例化为eval模式下的对象
auto model = torch::jit::load("resnet18.pt");
model.eval();
// 将模型迁移到GPU
model.to(torch::kCUDA);
// 构建随机输入张量并向其分配空间于指定位置(此处设为GPU)
const int batch_size = 1, channels = 3, height = 224, width = 224;
auto input_tensor = torch::randn({batch_size, channels, height, width}, torch::TensorOptions().device(torch::kCUDA));
// 执行前向传播计算得到输出结果
at::Tensor output = model.forward({input_tensor}).toTensor();
std::cout << "Output shape:" << output.sizes() << std::endl;
return 0;
}
```
此段代码展示了怎样通过调用API接口验证是否存在有效的CUDA资源;随后实例化一个预先保存好的神经网络架构(这里选用的是ResNet-18),将其转换成评估状态后迁移至图形处理单元之上进行高效运算。最后构造适当维度大小的数据集作为测试样本传递给该框架以获取预测类别概率分布情况。
阅读全文