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)运行,输出结果如下:
在这里插入图片描述

Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐