LibTorch实现MLP（多层感知机）

$$LibTorch实现MLP（多层感知机）$$

#include<opencv2/opencv.hpp>
#include <torch/torch.h>
#include <torch/script.h> 
using namespace std;

/// <summary>
/// 实现线性激活层
/// </summary>
class LinearReluImpl :public torch::nn::Module {
public:
    // 构造函数
	LinearReluImpl(int input, int output);
	// 前向传播
	torch::Tensor forward(torch::Tensor x);
private:
	// 线性层
	torch::nn::Linear linear1{ nullptr };
};

/// <summary>
/// 构造器，类外实现
/// </summary>
LinearReluImpl::LinearReluImpl(int input,int output) {
	// 将线性模板注册进参数组里，来给优化器更新权重参数
	linear1 = register_module("linear1", torch::nn::Linear(torch::nn::LinearOptions(input, output)));
}

/// <summary>
/// 前向传播，类外实现
/// </summary>
torch::Tensor LinearReluImpl::forward(torch::Tensor x) {
	// relu激活
	x = torch::relu(linear1(x));
	return x;
}


/// <summary>
/// 定义线性激活层xxx（相当于起别名）(和xxxImpl配套使用)
/// </summary>
/// <param name=""></param>
TORCH_MODULE(LinearRelu);


/// <summary>
/// 多层感知机类
/// </summary>
class MLP :public torch::nn::Module {
public:
	// 构造器
	MLP(int in_put,int out_put);
	// 前向传播
	torch::Tensor forward(torch::Tensor x);

private:
	// 中间层输入输出shape
	int mid_features[3] = { 32,64,128 };
	// 线性层
	LinearRelu ln1{ nullptr};
	LinearRelu ln2{ nullptr };
	LinearRelu ln3{ nullptr };
	// 最后输出层
	torch::nn::Linear out_ln{ nullptr };

};

/// <summary>
/// MLP构造器
/// </summary>
/// <returns></returns>
MLP::MLP(int in_features, int out_features) {
	ln1 = LinearRelu(in_features, mid_features[0]);
	ln2 = LinearRelu(mid_features[0], mid_features[1]);
	ln3 = LinearRelu(mid_features[1], mid_features[2]);
	out_ln = torch::nn::Linear(mid_features[2], out_features);
	// 注册模块，以优化器更新权重
	ln1 = register_module("ln1", ln1);
	ln2 = register_module("ln2", ln2);
	ln3 = register_module("ln3", ln3);
	out_ln = register_module("out_ln", out_ln);
}

/// <summary>
/// MLP前向传播
/// </summary>
/// <returns></returns>
torch::Tensor MLP::forward(torch::Tensor x) {
	x = ln1->forward(x);
	x = ln2->forward(x);
	x = ln3->forward(x);
	x = out_ln->forward(x);
	return x;
}


int main()
{

	//step0:定义使用cuda
	auto device = torch::Device(torch::kCUDA, 0);

	// step1:生成测试数据
	auto input = torch::ones({100});
	cout << input.sizes() << "  "<< input<<endl ;

	// step2:生成网络层实例
	auto net = MLP(100,10);
	
	// step3:推理输出
	try
	{
		auto output = net.forward(input);
		// step4:打印输出和大小
		cout << output << endl;
	}
	catch (const std::exception& e)
	{
		// step5:打印报错
		cout << e.what() << endl;
	}
	
	return 0;
}