深度学习与C++：从PyTorch到底层实现的映射关系

作者

在深度学习领域，PyTorch 以其灵活的动态计算图和易用的 Python 接口受到广泛关注。然而，对于需要极致性能和可定制化的场景，往往需要直接触碰 C++ 后端，了解其实现原理。本文将从几个核心维度——张量操作、自动求导、模块化设计以及性能优化，剖析 PyTorch C++ 后端（LibTorch）的实现思路，并展示如何在 C++ 中实现一个简易的线性层及其梯度计算。

1. 张量（Tensor）在 C++ 中的实现

PyTorch 的张量是 torch::Tensor，其内部结构由 TensorImpl 负责维护。关键点包括：

  • 内存布局：采用 N维连续内存，支持 C-order（行主序）和 F-order（列主序）。内存管理由 Allocator 抽象，默认使用 MallocAllocator，但也支持自定义内存池。
  • 数据类型：支持 torch::kFloat32, torch::kInt64 等多种类型。数据类型通过 TensorImpl 的 dtype() 获得。
  • 梯度追踪：Tensor 拥有 requires_grad 标志，并维护 grad_fn（梯度函数）来构建计算图。

在 C++ 里创建张量可以写成： 

torch::Tensor a = torch::randn({3, 3}, torch::requires_grad(true));

2. 自动求导（Autograd）机制

自动求导是 PyTorch 核心特性之一。其实现基于 Function 与 Backward 的双向遍历：

  1. Function：每个运算对应一个 torch::autograd::Function 子类，实现 forward 与 backward。
  2. 计算图：Tensor 的 grad_fn 指向其上一次运算的 Function，形成链式结构。
  3. 梯度回传：调用 tensor.backward() 时，系统从叶子节点向上遍历调用 backward，累计梯度。

示例：实现一个自定义的加法 Function。 

struct AddFunction : torch::autograd::Function <AddFunction> {
    static torch::Tensor forward(torch::autograd::AutogradContext *ctx,
                                 const torch::Tensor &a,
                                 const torch::Tensor &b) {
        ctx->save_for_backward({a, b});
        return a + b;
    }
    static std::vector<torch::Tensor> backward(torch::autograd::AutogradContext *ctx,
                                               std::vector<torch::Tensor> grad_outputs) {
        auto saved = ctx->get_saved_variables();
        return {grad_outputs[0], grad_outputs[0]};
    }
};

torch::Tensor add(const torch::Tensor &a, const torch::Tensor &b) {
    return AddFunction::apply(a, b);
}

3. 模块化设计（Modules）

PyTorch 的 nn.Module 在 C++ 里对应 torch::nn::Module，其核心思路是：

  • 子模块注册：register_module(name, module) 用于管理子模块，形成层级结构。
  • 参数注册：register_parameter(name, param) 用于注册可训练参数。
  • 前向传播：子类实现 forward() 方法。

3.1 简易线性层实现

struct LinearImpl : torch::nn::Module {
    torch::Tensor weight, bias;

    LinearImpl(int64_t in_features, int64_t out_features) {
        weight = register_parameter("weight",
            torch::empty({out_features, in_features}).normal_(0, 0.02));
        bias = register_parameter("bias",
            torch::zeros(out_features));
    }

    torch::Tensor forward(const torch::Tensor &x) {
        return torch::mm(x, weight.t()) + bias;
    }
};
TORCH_MODULE(Linear);  // 生成 Linear 类

3.2 计算梯度

在 C++ 里训练一个简单线性回归模型： 

int main() {
    // 数据
    torch::Tensor X = torch::randn({10, 3}, torch::requires_grad(false));
    torch::Tensor y = torch::randn({10, 1});

    // 模型
    Linear model(3, 1);
    torch::optim::SGD optim(model->parameters(), 0.01);

    for (int epoch = 0; epoch < 100; ++epoch) {
        optim.zero_grad();
        auto pred = model->forward(X);
        auto loss = torch::mse_loss(pred, y);
        loss.backward();
        optim.step();

        if (epoch % 10 == 0)
            std::cout << "epoch " << epoch << ", loss: " << loss.item<float>() << std::endl;
    }
}

4. 性能优化技巧

C++ 版本的 PyTorch 允许我们细粒度地控制性能：

  • 内存复用：使用 at::TensorOptions 指定 device、dtype，并配合 at::Allocator 的内存池。
  • JIT + TorchScript：将 C++ 模型通过 torch::jit::script::Module 编译，获得更快的运行时。
  • 多线程：开启 OMP_NUM_THREADS 或者使用 at::parallel_for 分块计算。
  • 显式指针管理：在需要时使用 torch::Tensor::data_ptr() 直接操作底层数据，减少拷贝。

5. 小结

通过本文的演示，我们了解了 PyTorch C++ 后端的核心实现机制：张量内存管理、自动求导、模块化设计与性能优化。掌握这些底层细节后，开发者可以在 C++ 环境中构建高效、可定制的深度学习模型，为嵌入式、游戏或实时系统提供强大支持。

祝你在 C++ 与深度学习的交叉道路上一帆风顺！

发表评论