libtorch初体验

it2023-06-06  79

0. 准备:

win10 VS2017 pytorch 1.1.0(cuda10.0) libtorch 1.1.0(cuda10.0) cmake >=3.12

解压libtorch放到自己喜欢的位置

1. pytorch简单实现一个MLP:

mlp.py

from torch import nn class MLP(nn.Module): def __init__(self, in_dim, hid_dim1, hid_dim2, out_dim): super(MLP, self).__init__() # 通过Sequential快速搭建三层的感知机 self.layer = nn.Sequential( nn.Linear(in_dim, hid_dim1), nn.ReLU(), nn.Linear(hid_dim1, hid_dim2), nn.ReLU(), nn.Linear(hid_dim2, out_dim), nn.ReLU()) def forward(self, x): x = self.layer(x) return x

2. 保存模型:

save.py

import torch from mlp import MLP model = MLP(28 * 28, 300, 200, 10) input = torch.ones((1, 28*28)) traced_script_module = torch.jit.trace(model, input) traced_script_module.save("model.pt") output = traced_script_module(input) print(output)

3. cpp工程:

1、 新建一个工程目录,并进入 2、 新建一个build文件夹 3、 新建main.cpp,修改内容如下:

#include <torch/script.h> // One-stop header. #include <iostream> #include <memory> int main(int argc, const char* argv[]) { torch::Device device = torch::kCUDA;//cpu的话改成torch::kCPU // Deserialize the ScriptModule from a file using torch::jit::load(). std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("D://Program//my_mlp//model.pt"); module->to(device); assert(module != nullptr); std::cout << "model load ok"<<std::endl; // Create a vector of inputs. std::vector<torch::jit::IValue> inputs; inputs.push_back(torch::ones({ 1, 28*28 }, device)); // Execute the model and turn its output into a tensor. at::Tensor output = module->forward(inputs).toTensor(); std::cout << output << std::endl; }

3、 新建CMakeLists.txt 4、 打开CMakeLists.txt,修改内容如下:

cmake_minimum_required(VERSION 3.12 FATAL_ERROR) project(Test) find_package(Torch REQUIRED) if(NOT Torch_FOUND) message(FATAL_ERROR "Pytorch Not Found!") endif(NOT Torch_FOUND) message(STATUS "Pytorch status:") message(STATUS "libraries: ${TORCH_LIBRARIES}") add_executable(Test main.cpp) target_link_libraries(Test ${TORCH_LIBRARIES}) set_property(TARGET Test PROPERTY CXX_STANDARD 11)

5、进入build文件夹打开cmd,输入如下并回车

cmake -DCMAKE_PREFIX_PATH=D:\ENVs\libtorch-DCMAKE_BUILD_TYPE=Release -G"Visual Studio 15 Win64" ..

其中D:\ENVs\libtorch应该替换为解压完libtorch的位置的下一级,我这里是D:\ENVs\libtorch(本文使用release版)

得到返回结果为如下,即ok

-- Configuring done -- Generating done -- Build files have been written to: D:/Program/libtorch_prog/test/build

4. VS运行:

1、 在build目录打开Test.sln

2、改为Release x64

3、 右键Test—>设为启动项目 3、 F5运行报错找不到dll 我把之前解压的libtorch\lib目录下的所有dll文件,复制到cpp工程下的build\Release文件夹中了

4、 再运行结果ok和save.py的输出一致

最新回复(0)