使用LibTorch训练模型(pytorch c++)

最新推荐文章于 2025-09-24 10:35:04 发布

原创

最新推荐文章于 2025-09-24 10:35:04 发布 · 4.7k 阅读

20 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #c++ #深度学习 #libtorch

本文介绍了如何使用libtorch进行C++14编程，并提供了LeNet5神经网络的实现步骤，包括编译设置、Hello World示例和创建自定义Dataset。同时涵盖了定义模型、数据预处理和训练流程的关键部分。

libtorch需要使用c++14的编译器

下载LibTorch，官网下载 https://pytorch.org/

调试的时候建议使用cpu+debug版本，到实际使用部署的时候再切换为cuda版本。

helloworld

#include "torch/library.h"
#include "torch/script.h"

int main()
{
	torch::Tensor output = torch::randn({ 3,2 });
	std::cout << output;
	
	return 0;
}

output

-1.9173 -0.5073
 1.5614 -0.0566
-0.0884  0.9237
[ CPUFloatType{3,2} ]

定义module

class LeNet5 : public torch::nn::Module
{
public:
	torch::nn::Conv2d C1;
	torch::nn::Conv2d C3;
	torch::nn::Linear F5;
	torch::nn::Linear F6;
	torch::nn::Linear OUTPUT;

public:
	LeNet5() 
		:
		// (inputchannel, outputchannel, kernel_size)
		C1(torch::nn::Conv2d(torch::nn::Conv2dOptions(3, 6, 5).padding(2).padding_mode(torch::kCircular))),
		C3(torch::nn::Conv2d(6, 16, 5)),
		F5(torch::nn::Linear(16 * 5 * 5, 120))