libtorch train example --- lenet5
libtorch是pytorch的C++版深度学习平台,可以通过C++代码完成深度学习任务中的常见操作。这里介绍如何构通过libtorch构建一个demo程序—lenet5 。源代码https://github.com/lishiyu93/libtorch_lenet5_train
libtorch是pytorch的C++版深度学习平台,可以通过C++代码完成深度学习任务中的常见操作。这里介绍如何构通过libtorch构建一个demo程序—lenet5 。
资源
源代码
https://github.com/lishiyu93/libtorch_lenet5_train
libtorch1.6_Debug版下载
libtorch1.6_Release版下载
环境
硬件:
CPU:Intel®Core™i5-10400F CPU @2.90GHz 2.90GHz
GPU: NVIDIA GeForce GTX 1660 Ti (6G显存)
内存:16.0GB(15.9GB可用)
软件:
Visual Studio 2019
Cmake
CUDA10.0/10.1,cudnn
OpenCV 3.4.0
libtorch1.6
CMake构建工程
CMakeLists修改
cmake_minimum_required(VERSION 3.0)
project(lenet5) #1
set(OpenCV_DIR "D:\\soft\\Opencv\\opencv\\build") #2
set(OpenCV_INCLUDE_DIRS ${OpenCV_DIR}\\include) #3
set(OpenCV_LIB_DIRS ${OpenCV_DIR}\\x64\\vc15\\lib) #4
set(OpenCV_Debug_LIBS "opencv_world340d.lib") #5
set(OpenCV_Release_LIBS "opencv_world340.lib") #6
set(Libtorch_Debug_DIR "D:\\Git_Project\\libtorch_train\\libtorch") #7
set(Libtorch_Debug_INCLUDE_DIRS ${Libtorch_Debug_DIR}\\include) #8
set(Libtorch_Debug_LIB_DIRs ${Libtorch_Debug_DIR}\\lib) #9
#setup CUDA #10
find_package(CUDA REQUIRED)
message(STATUS " libraries: ${CUDA_LIBRARIES}")
message(STATUS " include path: ${CUDA_INCLUDE_DIRS}")
include_directories(${CUDA_INCLUDE_DIRS})
enable_language(CUDA) # add this line, then no need to setup cuda path in vs
#setup OpenCV #11
include_directories(${OpenCV_INCLUDE_DIRS})
link_directories(${OpenCV_LIB_DIRS})
#setup libtorch #12
include_directories(${Libtorch_Debug_INCLUDE_DIRS})
include_directories(${Libtorch_Debug_INCLUDE_DIRS}\\torch\\csrc\\api\\include)
link_directories(${Libtorch_Debug_LIB_DIRs})
add_executable(lenet5 ${PROJECT_SOURCE_DIR}/main.cpp) #13
#setup additional dependencies #14
target_link_libraries(lenet5 ${CUDA_LIBRARIES}) #15
target_link_libraries(lenet5 debug ${OpenCV_Debug_LIBS}) #16
target_link_libraries(lenet5 optimized ${OpenCV_Release_LIBS}) #17
target_link_libraries(lenet5 "caffe2_nvrtc.lib" "c10.lib" "c10_cuda.lib") #18
target_link_libraries(lenet5 "torch.lib" "torch_cuda.lib" "torch_cpu.lib") #19
CMakeList,具体内容如下:
1)设置工程名:#1
2)设置OpenCV路径:#2—#6
3)设置Debug版libtorch库路径:#7—#9
4)设置CUDA路径:#10
5)配置OpenCV包含目录和库目录:#11
6)配置libtorch包含目录和库目录:#12
7)添加工程文件:#13
8)配置CUDA的附加依赖项:#15
9)配置OpenCV的debug版附加依赖项:#16
10)配置OpenCV的release版附加依赖项:#17
11)配置OpenCV的debug版附加依赖项:#18,#19
构建工程
训练代码
#include <iostream>
#include <chrono>
#include <time.h>
#include <iomanip>
#include <torch/torch.h>
#include <opencv2/opencv.hpp>
#include <opencv2/imgproc/imgproc.hpp>
using namespace std;
//using namespace std::chrono;
struct Net : torch::nn::Module {
Net()
: conv1(torch::nn::Conv2dOptions(1, 6, /*kernel_size=*/5).stride(1).padding(0)),
pool1(torch::nn::AvgPool2dOptions({ 2, 2 }).stride({ 2, 2 }).padding(0)),
conv2(torch::nn::Conv2dOptions(6, 16, /*kernel_size=*/5).stride(1).padding(0)),
fc1(16 * 5 * 5, 120),
fc2(120, 84),
fc3(84, 10) {
register_module("conv1", conv1);
register_module("pool1", pool1);
register_module("conv2", conv2);
register_module("fc1", fc1);
register_module("fc2", fc2);
register_module("fc3", fc3);
}
torch::Tensor forward(torch::Tensor x) {
cout << "input: " << x.size(0) << " " << x.size(1) << " " << x.size(2) << " " << x.size(3) << endl;
x = torch::relu(conv1->forward(x));
cout << "conv1: " << x.size(0) << " " << x.size(1) << " " << x.size(2) << " " << x.size(3) << endl;
x = pool1->forward(x);
cout << "pool1: " << x.size(0) << " " << x.size(1) << " " << x.size(2) << " " << x.size(3) << endl;
x = torch::relu(conv2->forward(x));
cout << "conv2: " << x.size(0) << " " << x.size(1) << " " << x.size(2) << " " << x.size(3) << endl;
x = pool1->forward(x);
cout << "pool2: " << x.size(0) << " " << x.size(1) << " " << x.size(2) << " " << x.size(3) << endl;
x = x.view({ x.size(0), -1 });
cout << "view: " << x.size(1) << endl;
x = torch::relu(fc1->forward(x));
cout << "fc1: " << x.size(1) << endl;
x = torch::relu(fc2->forward(x));
x = torch::softmax(fc3->forward(x), /*dim=*/1);
return x;
}
torch::nn::Conv2d conv1;
torch::nn::AvgPool2d pool1;
torch::nn::Conv2d conv2;
torch::nn::Linear fc1;
torch::nn::Linear fc2;
torch::nn::Linear fc3;
};
int
main(int argc,
char** argv)
{
// Device
torch::DeviceType device_type;
if (torch::cuda::is_available())
{
device_type = torch::kCUDA;
}
else
{
device_type = torch::kCPU;
}
torch::Device device(device_type);
torch::manual_seed(1234);
auto net = std::make_shared<Net>();
net->to(device);
net->eval();
torch::Tensor tmp = torch::ones({1, 1, 32, 32}).to(device);
net->train();
torch::Tensor out = net->forward(tmp);
std::cout << "lenet out shape: " << out.size(1) << std::endl;
std::cout << "lenet out: " << out << std::endl;
torch::save(net, "lenet5_c.pt");
std::cout << "train done." << std::endl;
return (0);
}
测试
1)编译工程
2)将/libtorch/lib/下所有dll和lib复制到/build/Debug/目录下,不然回报错缺少dll,例如:
3)运行,输出结果如下:
开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!
更多推荐
所有评论(0)