本文共 2074 字,大约阅读时间需要 6 分钟。
我们首先实现了一个 LinearRelu
类,该类主要包含一个线性层和一个 ReLU 激活函数。具体实现如下:
LinearReluImpl: public torch::nn::Module { public: LinearReluImpl(int input, int output); torch::Tensor forward(torch::Tensor x); private: torch::nn::Linear linear1; };
构造函数:
LinearReluImpl::LinearReluImpl(int input, int output) { linear1 = register_module("linear1", torch::nn::Linear(torch::nn::LinearOptions(input, output))); }
前向传播:
torch::Tensor LinearReluImpl::forward(torch::Tensor x) { x = torch::relu(linear1(x)); return x; }
接下来,我们定义了一个 MLP
类,它继承自 torch::nn::Module
。该类实现了一个三个隐藏层的MLP模型:
MLP: public torch::nn::Module { public: MLP(int in_features, int out_features); torch::Tensor forward(torch::Tensor x); private: int mid_features[3] = {32, 64, 128}; LinearRelu ln1, ln2, ln3; torch::nn::Linear out_ln; };
构造函数:
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); }
前向传播:
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; }
在main函数中,我们展示了如何使用我们的MLP模型:
int main() { // 检查是否使用 CUDA auto device = torch::Device(torch::kCUDA, 0); // 生成样本数据 auto input = torch::ones({100, 3}, device); // 初始化模型 MLP model(3, 10); // 前向传播 auto output = model(input); // 打印结果 std::cout << "输入大小: " << input.sizes() << std::endl; std::cout << "输出大小: " << output.sizes() << std::endl << std::endl; return 0; }
这样,我们完整地实现了一个使用 LibTorch 的多层感知机模型。整个实现过程包括定义激活函数和网络结构,并展示了如何在实际应用中使用这个模型。
转载地址:http://awwfk.baihongyu.com/