pytorch c++

#include <torch/torch.h>
#include <iostream>

// 定义一个简单的神经网络模型
struct Net : torch::nn::Module {
    Net() {
        // 定义两个线性层
        fc1 = register_module("fc1", torch::nn::Linear(784, 64));
        fc2 = register_module("fc2", torch::nn::Linear(64, 10));
    }

    // 前向传播函数
    torch::Tensor forward(torch::Tensor x) {
        // 将输入展平，并通过第一层线性层，然后应用ReLU激活函数
        x = torch::relu(fc1->forward(x.reshape({x.size(0), 784})));
        // 通过第二层线性层，并应用log_softmax
        x = torch::log_softmax(fc2->forward(x), 1);
        return x;
    }

    // 线性层成员变量
    torch::nn::Linear fc1{nullptr}, fc2{nullptr};
};

int main() {
    // 初始化网络
    Net net;

    // 创建一个随机输入张量 (模拟MNIST数据集的输入)
    auto input = torch::randn({64, 1, 28, 28});

    // 获取网络的前向传播输出
    auto output = net.forward(input);

    // 打印输出张量的形状
    std::cout << "Output shape: " << output.sizes() << std::endl;

    return 0;
}

解释说明

1. 引入头文件：

   • #include <torch/torch.h>：包含PyTorch C++库的所有必要定义。
   • #include <iostream>：用于标准输入输出。

2. 定义神经网络模型：

   • Net 类继承自 torch::nn::Module，这是 PyTorch 中所有模块的基类。
   • 在构造函数中，我们定义了两个线性层（fc1 和 fc2），并使用 register_module 方法将它们注册到网络中。

3. 前向传播函数：

   • forward 函数定义了网络的前向传播逻辑。它接收一个输入张量 x，首先将其展平为一维向量，然后通过第一层线性层和 ReLU 激活函数，最后通过第二层线性层并应用 log_softmax。

4. 主函数：

   • 初始化 Net 实例。
   • 创建一个随机输入张量，模拟 MNIST 数据集的输入（大小为 [batch_size, channels, height, width]）。
   • 调用 forward 函数获取输出，并打印输出张量的形状。

这个示例展示了如何使用 PyTorch C++ API 构建和运行一个简单的神经网络。