PyTorch C++ API笔记

最新推荐文章于 2026-06-23 11:12:53 发布

原创

最新推荐文章于 2026-06-23 11:12:53 发布 · 2.5k 阅读

标签

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

本文详细介绍了PyTorch的C++ API，包括ATen、autograd、C++前端、TorchScript和C++扩展。ATen是基础张量和数学运算库，autograd提供了自动微分功能。C++前端包含torch::Tensor、torch::nn等模块，支持自动微分和高级机器学习应用。TorchScript是一种编程语言，用于模型的序列化和编译。C++扩展则允许Python和C++之间的交互。

官方文档地址：
https://pytorch.org/cppdocs/

说明：以下是我学习笔记，大部分内容摘录自官方文档，也添加了一些案例。

PyTorch C++ API大致可以分为五个部分：

ATen：基础张量和数学运算库，所有其他功能都建立在它之上。
Autograd：通过自动微分增强 ATen。
C++前端：用于训练和评估机器学习模型的高级构造。
TorchScript：TorchScript JIT 编译器和解释器的接口。
C++扩展：一种使用自定义 C++ 和 CUDA 例程扩展 Python API 的方法。

结合起来，这些构建块形成了一个研究和生产就绪的 C++ 库，用于张量计算和动态神经网络，重点是 GPU 加速和快速 CPU 性能。

1.ATen

ATen 本质上是一个张量库，PyTorch 中几乎所有其他 Python 和 C++ 接口都建立在它之上。它提供了一个核心Tensor类，在该类上定义了数百个操作。

#include <ATen/ATen.h>

at::Tensor a = at::ones({
   
   2, 2}, at::kInt);
at::Tensor b = at::randn({
   
   2, 2});
auto c = a + b.to(at::kInt);

2.autograd

autograd是 PyTorch 的 C++ API 的一部分，它Tensor通过自动微分功能增强了 ATen类。autograd 系统记录对张量的操作以形成autograd 图。调用backwards()执行反向模式微分，最终产生梯度。

#include <torch/csrc/autograd/variable.h>
#include <torch/csrc/autograd/function.h>

torch::Tensor a = torch::ones({
   
   2, 2}, torch::requires_grad());
torch::Tensor b = torch::randn({
   
   2, 2});
auto c = a + b;
c.backward(); // a.grad() will now hold the gradient of c w.r.t. a.

注意：这里需要关注at::Tensor和torch::Tensor的区别！
ATen 中的类默认是不可微的。要添加 autograd API 提供的Tensor的可微性，就必须使用来自torch::命名空间，而不是at::命名空间的Tensor函数。例如，虽然使用at::ones创建的Tensor将不可微，但使用torch::ones创建的Tensor将是可微的。