libtorch train example --- lenet5

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）运行，输出结果如下：