目录

使用 PyTorch C++ 前端

创建日期: 2019年1月15日 | 最后更新日期: 2024年12月18日 | 最后验证日期: 2024年11月5日

PyTorch的C++前端是一个纯C++接口,用于访问PyTorch机器学习框架。虽然PyTorch的主要接口自然是Python,但这个Python API建立在庞大的C++代码基础上,提供诸如张量和自动微分等基础数据结构和功能。C++前端暴露了一个纯C++11 API,该API在此底层C++代码基础上扩展了用于机器学习训练和推理所需的工具。这包括内置的常见神经网络建模组件集合;一个自定义模块的API,可以扩展此组件集合;一系列流行的优化算法库,如随机梯度下降;并行数据加载器及其API,用于定义和加载数据集;序列化例行程序等。

本教程将引导您完成使用C++前端训练模型的端到端示例。具体来说,我们将训练一个DCGAN——一种生成模型——以生成MNIST数字图像。虽然从概念上讲这是一个简单的示例,但它应该足以让您对PyTorch C++前端有一个快速了解,并激发您训练更复杂模型的兴趣。我们首先会介绍一些使用C++前端的原因,然后直接进入模型的定义和训练。

提示

观看 这个来自CppCon 2018的闪电演讲,快速(且幽默地)了解C++前端的介绍。

提示

本说明提供了C++前端组件和设计哲学的全面概述。

提示

Documentation for the PyTorch C++ 生态系统可在 https://pytorch.org/cppdocs 获取。在那里,您可以找到高层次描述以及API级别的文档。

动机

在我们开始令人兴奋的生成对抗网络(GANs)和MNIST数字的研究之旅之前,让我们退一步,讨论一下为什么你会选择从一开始就使用C++前端而不是Python前端。我们(PyTorch团队)创建了C++前端,目的是在Python无法使用或者并不是最适合的环境中进行研究。这样的环境包括:

  • 低延迟系统: 您可能希望在具有高帧率和低延迟要求的纯C++游戏引擎中进行强化学习研究。在这种环境中,使用纯C++库比Python库更合适。由于Python解释器的速度较慢,Python可能根本不可行。

  • 高度多线程环境:由于全局解释器锁 (GIL),Python 无法同时运行多个系统线程。多进程是另一种选择,但不如多线程可扩展且存在显著的缺点。C++ 没有这样的限制,线程易于使用和创建。需要大量并行化的模型,如用于 深度神经进化 的模型,可以从这一点中受益。

  • 现有C++代码库: 您可能已经是某个现有的C++应用程序的所有者,该应用程序从后端服务器提供网页服务,或者在照片编辑软件中渲染3D图形,并希望将机器学习方法集成到您的系统中。C++前端允许您继续使用C++,避免在Python和C++之间来回绑定的麻烦,同时保留传统PyTorch(Python)体验中的大部分灵活性和直观性。

C++前端并不是要与Python前端竞争,而是为了补充它。我们知道研究人员和工程师都喜欢PyTorch的简洁性、灵活性以及直观的API。我们的目标是在所有可能的环境中确保您能够利用这些核心设计原则,包括上述描述的各种环境。如果您的情况符合其中一种场景,或者您只是单纯地感兴趣或好奇,请继续阅读接下来几段,我们将详细探讨C++前端。

提示

C++前端尽量提供与Python前端尽可能接近的API。如果你对Python前端很熟悉,并且有时会问自己“用C++前端如何实现X?”那么你可以按照在Python中编写代码的方式进行编写,通常情况下,C++中也会有与Python相同的函数和方法(只需记得用双冒号替换点号即可)。

编写基本应用程序

让我们首先编写一个最小的C++应用程序,以验证我们的设置和构建环境是否一致。首先,您需要获取LibTorch发行版的副本——我们的预构建zip归档文件,其中包含了使用C++前端所需的所有相关头文件、库和CMake构建文件。LibTorch发行版可以在PyTorch网站上下载,适用于Linux、MacOS和Windows。本教程的其余部分将假设您使用的是基本的Ubuntu Linux环境,但您也可以在MacOS或Windows上进行操作。

提示

关于安装 PyTorch 的 C++ 发行版的说明详细描述了以下步骤。

提示

在Windows上,调试版和发布版的ABI不兼容。如果您计划以调试模式构建项目,请尝试使用LibTorch的调试版本。同时,请确保在下面的第cmake --build .行中指定了正确的配置。

第一步是通过从PyTorch官网获取的链接,在本地下载LibTorch分发包。对于纯正的Ubuntu Linux环境,这需要运行以下命令:

# If you need e.g. CUDA 9.0 support, please replace "cpu" with "cu90" in the URL below.
wget https://download.pytorch.org/libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zip
unzip libtorch-shared-with-deps-latest.zip

接下来,让我们编写一个名为dcgan.cpp的小型C++文件,该文件包含 torch/torch.h 并且现在仅打印出一个3x3单位矩阵:

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

int main() {
  torch::Tensor tensor = torch::eye(3);
  std::cout << tensor << std::endl;
}

为了构建这个小型应用程序以及我们稍后完整的训练脚本,我们将使用这个 CMakeLists.txt 文件:

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(dcgan)

find_package(Torch REQUIRED)

add_executable(dcgan dcgan.cpp)
target_link_libraries(dcgan "${TORCH_LIBRARIES}")
set_property(TARGET dcgan PROPERTY CXX_STANDARD 14)

注意

虽然 CMake 是推荐的构建系统,但它并非强制要求。你也可以使用 Visual Studio 项目文件、QMake、纯 Makefile 或任何其他你觉得舒适的构建环境。不过,我们并不提供开箱即用的支持。

注意上面 CMake 文件中的第 4 行: find_package(Torch REQUIRED). 这指示 CMake 查找 LibTorch 库的构建配置。 为了使 CMake 知道在哪里查找这些文件,我们在调用 cmake 时必须设置 CMAKE_PREFIX_PATH。在我们进行此操作之前,让我们为我们的 dcgan 应用程序约定以下目录结构:

dcgan/
  CMakeLists.txt
  dcgan.cpp

进一步,我将把解压后的LibTorch分布路径称为/path/to/libtorch。请注意,这必须是绝对路径。特别是,将CMAKE_PREFIX_PATH设置为类似于../../libtorch的内容会导致意外的问题。相反,请使用$PWD/../../libtorch来获取相应的绝对路径。现在,我们准备好构建我们的应用程序:

root@fa350df05ecf:/home# mkdir build
root@fa350df05ecf:/home# cd build
root@fa350df05ecf:/home/build# cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
-- The C compiler identification is GNU 5.4.0
-- The CXX compiler identification is GNU 5.4.0
-- Check for working C compiler: /usr/bin/cc
-- Check for working C compiler: /usr/bin/cc -- works
-- Detecting C compiler ABI info
-- Detecting C compiler ABI info - done
-- Detecting C compile features
-- Detecting C compile features - done
-- Check for working CXX compiler: /usr/bin/c++
-- Check for working CXX compiler: /usr/bin/c++ -- works
-- Detecting CXX compiler ABI info
-- Detecting CXX compiler ABI info - done
-- Detecting CXX compile features
-- Detecting CXX compile features - done
-- Looking for pthread.h
-- Looking for pthread.h - found
-- Looking for pthread_create
-- Looking for pthread_create - not found
-- Looking for pthread_create in pthreads
-- Looking for pthread_create in pthreads - not found
-- Looking for pthread_create in pthread
-- Looking for pthread_create in pthread - found
-- Found Threads: TRUE
-- Found torch: /path/to/libtorch/lib/libtorch.so
-- Configuring done
-- Generating done
-- Build files have been written to: /home/build
root@fa350df05ecf:/home/build# cmake --build . --config Release
Scanning dependencies of target dcgan
[ 50%] Building CXX object CMakeFiles/dcgan.dir/dcgan.cpp.o
[100%] Linking CXX executable dcgan
[100%] Built target dcgan

上面,我们首先在dcgan目录中创建了一个build文件夹,进入这个文件夹后,运行了cmake命令来生成必要的构建(Make)文件,并最终通过运行cmake --build . --config Release成功编译了项目。现在我们已经准备好执行我们的最小二进制文件并完成这个关于基本项目配置的部分:

root@fa350df05ecf:/home/build# ./dcgan
1  0  0
0  1  0
0  0  1
[ Variable[CPUFloatType]{3,3} ]

看起来像一个单位矩阵!

定义神经网络模型

现在我们的基本环境已经配置好了,我们可以深入探讨这个教程中更加有趣的部分了。首先,我们将讨论如何在C++前端定义和交互使用模块。我们将会从基本的小规模示例模块开始,然后利用C++前端提供的丰富内置模块库实现一个完整的生成对抗网络(GAN)。

模块API基础

与Python接口一致,基于C++前端的神经网络由称为模块的可重用构建块组成。有一个基础模块类,所有其他模块均派生自该类。在Python中,该类是torch.nn.Module,在C++中它是torch::nn::Module。除了实现模块封装算法的forward()方法外,模块通常包含三种类型的子对象中的任意一种:参数、缓冲区和子模块。

参数和缓冲区以张量的形式存储状态。参数记录梯度,而缓冲区不记录。参数通常是您神经网络中的可训练权重。批归一化的均值和方差是缓冲区的例子。为了重用特定的逻辑和状态,PyTorch API 允许模块嵌套。嵌套模块称为 子模块

参数、缓冲区和子模块必须显式注册。注册后,可以使用类似parameters()buffers()的方法来检索整个(嵌套)模块层次结构中的所有参数容器。 同样,类似to(...)的方法,例如to(torch::kCUDA)将所有参数和缓冲区从CPU内存移动到CUDA内存,也适用于整个模块层次结构。

定义一个模块和注册参数

要把这些话变成代码,让我们考虑一下用 Python 接口编写的这个简单模块:

import torch

class Net(torch.nn.Module):
  def __init__(self, N, M):
    super(Net, self).__init__()
    self.W = torch.nn.Parameter(torch.randn(N, M))
    self.b = torch.nn.Parameter(torch.randn(M))

  def forward(self, input):
    return torch.addmm(self.b, input, self.W)

在C++中,它会像这样看起来:

#include <torch/torch.h>

struct Net : torch::nn::Module {
  Net(int64_t N, int64_t M) {
    W = register_parameter("W", torch::randn({N, M}));
    b = register_parameter("b", torch::randn(M));
  }
  torch::Tensor forward(torch::Tensor input) {
    return torch::addmm(b, input, W);
  }
  torch::Tensor W, b;
};

就像在Python中一样,我们定义一个名为Net的类(为了简化,这里使用struct而不是class)并从模块基类派生它。 在构造函数中,我们使用torch::randn创建张量,就像在Python中使用torch.randn一样。一个有趣的区别在于我们如何注册参数。在Python中,我们将张量包装在torch.nn.Parameter类中,而在C++中,我们必须通过register_parameter方法传递张量。这是因为Python API可以检测一个属性是否为torch.nn.Parameter类型,并自动注册此类张量。在C++中,反射功能非常有限,因此提供了一种更传统(也更不神秘)的方法。

注册子模块和遍历模块层次结构

同样地,我们可以注册参数,也可以注册子模块。在 Python 中,当子模块被赋值为模块的一个属性时,它们会自动被检测并注册:

class Net(torch.nn.Module):
  def __init__(self, N, M):
      super(Net, self).__init__()
      # Registered as a submodule behind the scenes
      self.linear = torch.nn.Linear(N, M)
      self.another_bias = torch.nn.Parameter(torch.rand(M))

  def forward(self, input):
    return self.linear(input) + self.another_bias

这允许,例如,使用 parameters() 方法递归访问我们模块层次结构中的所有参数:

>>> net = Net(4, 5)
>>> print(list(net.parameters()))
[Parameter containing:
tensor([0.0808, 0.8613, 0.2017, 0.5206, 0.5353], requires_grad=True), Parameter containing:
tensor([[-0.3740, -0.0976, -0.4786, -0.4928],
        [-0.1434,  0.4713,  0.1735, -0.3293],
        [-0.3467, -0.3858,  0.1980,  0.1986],
        [-0.1975,  0.4278, -0.1831, -0.2709],
        [ 0.3730,  0.4307,  0.3236, -0.0629]], requires_grad=True), Parameter containing:
tensor([ 0.2038,  0.4638, -0.2023,  0.1230, -0.0516], requires_grad=True)]

要使用C++注册子模块,请使用恰当地命名的 register_module() 方法 来注册一个类似于 torch::nn::Linear 的模块:

struct Net : torch::nn::Module {
  Net(int64_t N, int64_t M)
      : linear(register_module("linear", torch::nn::Linear(N, M))) {
    another_bias = register_parameter("b", torch::randn(M));
  }
  torch::Tensor forward(torch::Tensor input) {
    return linear(input) + another_bias;
  }
  torch::nn::Linear linear;
  torch::Tensor another_bias;
};

提示

您可以在此处找到可用的内置模块完整列表,如 torch::nn::Lineartorch::nn::Dropouttorch::nn::Conv2d,在 torch::nn 命名空间的文档中这里

上述代码的一个细微之处在于为什么子模块在构造函数的初始化列表中创建,而参数则在构造函数主体中创建。这背后有一个很好的原因,我们将在下面更详细的C++前端的所有权模型部分提及。然而,最终的结果是我们可以像在Python中一样递归访问模块树的参数。调用parameters()返回一个std::vector<torch::Tensor>,我们可以对其进行迭代:

int main() {
  Net net(4, 5);
  for (const auto& p : net.parameters()) {
    std::cout << p << std::endl;
  }
}

打印的内容为:

root@fa350df05ecf:/home/build# ./dcgan
0.0345
1.4456
-0.6313
-0.3585
-0.4008
[ Variable[CPUFloatType]{5} ]
-0.1647  0.2891  0.0527 -0.0354
0.3084  0.2025  0.0343  0.1824
-0.4630 -0.2862  0.2500 -0.0420
0.3679 -0.1482 -0.0460  0.1967
0.2132 -0.1992  0.4257  0.0739
[ Variable[CPUFloatType]{5,4} ]
0.01 *
3.6861
-10.1166
-45.0333
7.9983
-20.0705
[ Variable[CPUFloatType]{5} ]

使用三个参数,就像在Python中一样。要查看这些参数的名称,C++ API 提供了一个 named_parameters() 方法,返回一个 OrderedDict 类似于在Python中:

Net net(4, 5);
for (const auto& pair : net.named_parameters()) {
  std::cout << pair.key() << ": " << pair.value() << std::endl;
}

我们再次执行以查看输出:

root@fa350df05ecf:/home/build# make && ./dcgan                                                                                                                                            11:13:48
Scanning dependencies of target dcgan
[ 50%] Building CXX object CMakeFiles/dcgan.dir/dcgan.cpp.o
[100%] Linking CXX executable dcgan
[100%] Built target dcgan
b: -0.1863
-0.8611
-0.1228
1.3269
0.9858
[ Variable[CPUFloatType]{5} ]
linear.weight:  0.0339  0.2484  0.2035 -0.2103
-0.0715 -0.2975 -0.4350 -0.1878
-0.3616  0.1050 -0.4982  0.0335
-0.1605  0.4963  0.4099 -0.2883
0.1818 -0.3447 -0.1501 -0.0215
[ Variable[CPUFloatType]{5,4} ]
linear.bias: -0.0250
0.0408
0.3756
-0.2149
-0.3636
[ Variable[CPUFloatType]{5} ]

注意

文档 for torch::nn::Module 包含了操作模块层次结构的所有方法的完整列表。

运行网络的前向模式

要使用C++执行网络,我们只需调用自己定义的 forward() 方法:

int main() {
  Net net(4, 5);
  std::cout << net.forward(torch::ones({2, 4})) << std::endl;
}

这会打印类似的内容:

root@fa350df05ecf:/home/build# ./dcgan
0.8559  1.1572  2.1069 -0.1247  0.8060
0.8559  1.1572  2.1069 -0.1247  0.8060
[ Variable[CPUFloatType]{2,5} ]

模块所有权

在这一点上,我们知道如何在C++中定义一个模块,注册参数, 注册子模块,通过诸如parameters()的方法遍历模块层次结构,最后运行模块的forward()方法。虽然 C++ API 中还有许多其他方法、类和主题需要学习,但我将您推荐到文档以获取完整的内容。我们还将在接下来的实现过程中 介绍一些更多的概念,例如 DCGAN 模型和端到端训练管道。在我们这样做之前, 让我简要介绍一下 C++ 前端为 torch::nn::Module 的子类提供的所有权模型

对于本次讨论,所有权模型指的是模块的存储和传递方式——这决定了某个特定模块实例的所有权归属。在Python中,对象始终动态分配(在堆上)并且具有引用语义。这种方式非常易于使用且容易理解。实际上,在Python中,你可以几乎完全忽略对象的位置和引用方式,专注于完成任务。

C++, 作为一种低级语言,提供了更多的选项。这增加了复杂性,并且严重影响了 C++ 前端的设计和易用性。特别是对于 C++ 前端中的模块,我们有使用 值语义引用语义 的选择。第一种情况是最简单的,之前的示例已经展示了这一点:模块对象在栈上分配,当传递给函数时,可以被复制、移动(使用 std::move)或通过引用或指针获取:

struct Net : torch::nn::Module { };

void a(Net net) { }
void b(Net& net) { }
void c(Net* net) { }

int main() {
  Net net;
  a(net);
  a(std::move(net));
  b(net);
  c(&net);
}

对于第二种情况——引用语义,我们可以使用 std::shared_ptr。 引用语义的优势在于,就像在 Python 中一样,它减少了思考如何将模块传递给函数以及如何声明参数的认知负担(假设你在 everywhere 使用 shared_ptr)。

struct Net : torch::nn::Module {};

void a(std::shared_ptr<Net> net) { }

int main() {
  auto net = std::make_shared<Net>();
  a(net);
}

在我们的经验中,来自动态语言的研究人员更偏好引用语义而非值语义,尽管后者更“原生”地属于C++。同样重要的是要注意,torch::nn::Module的设计,为了接近Python API的使用体验,依赖于共享所有权。例如,看看我们之前(这里简化了)对Net的定义:

struct Net : torch::nn::Module {
  Net(int64_t N, int64_t M)
    : linear(register_module("linear", torch::nn::Linear(N, M)))
  { }
  torch::nn::Linear linear;
};

为了使用linear子模块,我们希望直接将其存储在我们的类中。然而,我们也希望模块基类能够了解并访问这个子模块。为此,它必须存储一个对该子模块的引用。此时,我们已经意识到共享所有权的需求。无论是torch::nn::Module类还是具体的Net类都需要对该子模块进行引用。因此,基类将模块存储为shared_ptr,具体类也必须如此。

但是等一下!我在上面的代码中没有看到任何关于shared_ptr的提及!为什么会这样? 其实是因为std::shared_ptr<MyModule>输入起来非常繁琐。为了保持研究人员的高效工作,我们设计了一种复杂的方案来隐藏对shared_ptr的提及——这通常是值语义的一个优点——同时保留引用语义。要理解这是如何工作的,我们可以看一下核心库中torch::nn::Linear模块的简化定义(完整的定义请参见这里):

struct LinearImpl : torch::nn::Module {
  LinearImpl(int64_t in, int64_t out);

  Tensor forward(const Tensor& input);

  Tensor weight, bias;
};

TORCH_MODULE(Linear);

在简要说明中:模块并非称为Linear,而是称为LinearImpl。然后,一个宏TORCH_MODULE定义了实际的Linear类。这个“生成”的类实际上是某个std::shared_ptr<LinearImpl>的包装器。它是一个包装器而不是简单的typedef,这样可以确保构造函数仍然按预期工作,即你可以仍然写torch::nn::Linear(3, 4)而不是std::make_shared<LinearImpl>(3, 4)。我们称宏创建的类为模块持有者。就像使用(共享)指针一样,你可以通过箭头运算符访问底层对象(例如model->forward(...))。最终的结果是一种所有权模型,与Python API非常相似。引用语义成为默认设置,但无需额外输入std::shared_ptrstd::make_shared。对于我们的Net,使用模块持有者API如下所示:

struct NetImpl : torch::nn::Module {};
TORCH_MODULE(Net);

void a(Net net) { }

int main() {
  Net net;
  a(net);
}

在这里值得一提的一个微妙问题是,一个默认构造的 std::shared_ptr 是“空”的,即包含空指针。什么是默认构造的 LinearNet? 嗯,这是一个棘手的选择。我们可以认为它应该是一个空的(null)std::shared_ptr<LinearImpl>。然而,请回忆一下 Linear(3, 4) 等同于 std::make_shared<LinearImpl>(3, 4)。这意味着如果我们决定 Linear linear; 应该是一个空指针,那么就无法构造一个不带任何构造函数参数的模块,或者将所有参数设为默认值。因此,在当前的API中,一个默认构造的模块持有者(如 Linear())会调用底层模块的默认构造函数(如 LinearImpl())。如果底层模块没有默认构造函数,将会得到编译错误。要构造空的持有者,可以将 nullptr 传递给持有者的构造函数。

在实践中,这意味着你可以像前面所示那样使用子模块,其中模块在初始化列表中注册并构建:

struct Net : torch::nn::Module {
  Net(int64_t N, int64_t M)
    : linear(register_module("linear", torch::nn::Linear(N, M)))
  { }
  torch::nn::Linear linear;
};

或者你可以先用空指针构造holder,然后在构造函数中为其赋值(对于Python开发者更为熟悉的方式):

struct Net : torch::nn::Module {
  Net(int64_t N, int64_t M) {
    linear = register_module("linear", torch::nn::Linear(N, M));
  }
  torch::nn::Linear linear{nullptr}; // construct an empty holder
};

总结:你应该使用哪种所有权模型——哪种语义?C++前端的API最适合支持模块持有者提供的所有权模型。这种机制唯一的缺点是在模块声明下方多了一行样板代码。不过,简单地说,C++模块介绍中展示的价值语义模型仍然是最简单的选择。对于小型、简单的脚本,你可能也能用它。但你会发现,由于技术原因,它并不总是被支持。例如,序列化API(torch::savetorch::load)只支持模块持有者(或普通的shared_ptr)。因此,模块持有者API是推荐的定义C++前端模块的方式,我们将在后续教程中使用这种方式。

定义DCGAN模块

我们现在有了必要的背景和介绍,可以定义我们在这篇文章中要解决的机器学习任务的模块。回顾一下:我们的任务是生成来自MNIST数据集的数字图像。我们希望使用生成对抗网络(GAN)来解决这个任务。特别是,我们将使用一种DCGAN架构——这是最早也是最简单的架构之一,但对于这个任务来说已经完全足够。

提示

您可以在此仓库中找到本教程中展示的完整源代码。

什么是GAN?

一个生成对抗网络(GAN)由两个独立的神经网络模型组成:一个 生成器 和一个 判别器。生成器接收来自噪声分布的样本,并试图将其转换为与目标分布(在我们的情况下是MNIST数据集)中的图像相似的图像。判别器则接收来自MNIST数据集的真实图像,或者来自生成器的虚假图像。它被要求发出一个概率判断某个特定图像的真实性(更接近于 1)或虚假性(更接近于 0)。生成器生成的图像的真实性反馈用于训练生成器。判别器对真实性判断能力的反馈用于优化判别器。理论上,生成器和判别器之间的微妙平衡使它们能够同步改进,最终导致生成器生成的图像与目标分布中的图像无法区分,使判别器(那时)出色的眼睛发出真实图像和虚假图像都具有相同概率(接近 0.5)的判断。对于我们来说,最终结果是一个机器,它接收噪声作为输入,并将其输出为真实的数字图像。

生成器模块

我们首先定义生成器模块,该模块由一系列转置的2D卷积、批量归一化和ReLU激活单元组成。我们在自定义的模块的forward()方法中显式地在模块之间传递输入(以函数方式传递):

struct DCGANGeneratorImpl : nn::Module {
  DCGANGeneratorImpl(int kNoiseSize)
      : conv1(nn::ConvTranspose2dOptions(kNoiseSize, 256, 4)
                  .bias(false)),
        batch_norm1(256),
        conv2(nn::ConvTranspose2dOptions(256, 128, 3)
                  .stride(2)
                  .padding(1)
                  .bias(false)),
        batch_norm2(128),
        conv3(nn::ConvTranspose2dOptions(128, 64, 4)
                  .stride(2)
                  .padding(1)
                  .bias(false)),
        batch_norm3(64),
        conv4(nn::ConvTranspose2dOptions(64, 1, 4)
                  .stride(2)
                  .padding(1)
                  .bias(false))
 {
   // register_module() is needed if we want to use the parameters() method later on
   register_module("conv1", conv1);
   register_module("conv2", conv2);
   register_module("conv3", conv3);
   register_module("conv4", conv4);
   register_module("batch_norm1", batch_norm1);
   register_module("batch_norm2", batch_norm2);
   register_module("batch_norm3", batch_norm3);
 }

 torch::Tensor forward(torch::Tensor x) {
   x = torch::relu(batch_norm1(conv1(x)));
   x = torch::relu(batch_norm2(conv2(x)));
   x = torch::relu(batch_norm3(conv3(x)));
   x = torch::tanh(conv4(x));
   return x;
 }

 nn::ConvTranspose2d conv1, conv2, conv3, conv4;
 nn::BatchNorm2d batch_norm1, batch_norm2, batch_norm3;
};
TORCH_MODULE(DCGANGenerator);

DCGANGenerator generator(kNoiseSize);

我们现在可以调用forward()DCGANGenerator上,将其噪声样本映射为图像。

选择的特定模块,如nn::ConvTranspose2dnn::BatchNorm2d, 遵循了之前概述的结构。kNoiseSize常量确定输入噪声向量的大小,并设置为100。超参数当然通过研究生下降法找到。

注意力

在发现超参数的过程中,并没有研究生受到伤害。他们只是定期食用了Soylent。

注意

A brief word on the way options are passed to built-in modules like Conv2d 在C++前端中内置模块(如Conv2d)的选项传递方式:每个模块都有某些必需的选项,例如BatchNorm2d模块所需的特征数量。如果你只需要配置必需的选项,可以直接将它们传递给模块的构造函数,例如BatchNorm2d(128)Dropout(0.5)Conv2d(8, 4, 2)(输入通道数、输出通道数和卷积核大小)。然而,如果你需要修改其他选项,这些选项通常会被默认设置,例如bias对于Conv2d,你需要构建并传递一个options对象。C++前端中的每个模块都关联有一个选项结构体,称为ModuleOptions,其中Module是模块的名称,例如LinearOptions对于Linear。这就是我们上面为Conv2d模块所做的。

判别模块

判别器同样是一系列卷积、批量归一化和激活函数。然而,这里的卷积现在是常规卷积而不是转置卷积,并且我们使用的是带alpha值为0.2的泄漏ReLU替代普通的ReLU。此外,最终的激活函数变为Sigmoid,它可以将值压缩到0和1之间。这样,我们可以将这些压缩后的值解释为判别器赋予图像真实性的概率。

要构建判别器,我们将尝试一些不同的方法:一个 Sequential 模块。 就像在 Python 中一样,这里 PyTorch 提供了两种模型定义的 API:一种是函数式 API,其中输入通过一系列函数传递(例如生成器模块示例);另一种是面向对象的 API,我们可以在其中构建一个包含整个模型的 Sequential 模块。使用 Sequential,判别器会看起来像这样:

nn::Sequential discriminator(
  // Layer 1
  nn::Conv2d(
      nn::Conv2dOptions(1, 64, 4).stride(2).padding(1).bias(false)),
  nn::LeakyReLU(nn::LeakyReLUOptions().negative_slope(0.2)),
  // Layer 2
  nn::Conv2d(
      nn::Conv2dOptions(64, 128, 4).stride(2).padding(1).bias(false)),
  nn::BatchNorm2d(128),
  nn::LeakyReLU(nn::LeakyReLUOptions().negative_slope(0.2)),
  // Layer 3
  nn::Conv2d(
      nn::Conv2dOptions(128, 256, 4).stride(2).padding(1).bias(false)),
  nn::BatchNorm2d(256),
  nn::LeakyReLU(nn::LeakyReLUOptions().negative_slope(0.2)),
  // Layer 4
  nn::Conv2d(
      nn::Conv2dOptions(256, 1, 3).stride(1).padding(0).bias(false)),
  nn::Sigmoid());

提示

一个 Sequential 模块简单地执行函数组合。第一个子模块的输出成为第二个子模块的输入,第三个子模块的输出成为第四个子模块的输入,依此类推。

加载数据

现在我们已经定义了生成器和判别器模型,需要一些可以用来训练这些模型的数据。C++ 前端,就像 Python 版本一样,配备了一个强大的并行数据加载器。这个数据加载器可以从一个数据集(你可以自己定义)中读取批量数据,并提供了许多配置选项。

注意

虽然Python数据加载器使用多进程,但C++数据加载器是真正的多线程,并不会启动新的进程。

数据加载器是C++前端的data api的一部分,位于torch::data::命名空间中。此API由几个不同的组件组成:

  • 数据加载器类,

  • 一个用于定义数据集的API,

  • 一个用于定义转换的API,可以应用于数据集,

  • 一个用于定义采样器的API,这些采样器生成用于索引数据集的索引,

  • 现有的数据集、转换和采样器库。

对于本教程,我们可以使用随附于C++前端的MNIST数据集。让我们为这个数据集实例化一个torch::data::datasets::MNIST,并应用两个转换:首先,我们将图像归一化到范围-1+1(原范围为01)。其次,我们将应用Stack 拼接,这会将一批张量沿第一个维度堆叠成一个单一的张量:

auto dataset = torch::data::datasets::MNIST("./mnist")
    .map(torch::data::transforms::Normalize<>(0.5, 0.5))
    .map(torch::data::transforms::Stack<>());

请注意,MNIST数据集应位于您执行训练二进制文件的相对路径下的./mnist目录中。您可以使用此脚本下载MNIST数据集。

接下来,我们创建一个数据加载器并将这个数据集传递给它。为了创建一个新的数据加载器,我们使用torch::data::make_data_loader,它返回的是正确类型的std::unique_ptr(这取决于数据集类型、采样器类型以及其他一些实现细节):

auto data_loader = torch::data::make_data_loader(std::move(dataset));

数据加载器确实有很多选项。你可以在这里查看完整的设置 这里。 例如,为了加快数据加载速度,我们可以增加工作线程的数量。默认数量是零,这意味着将使用主线程。 如果我们设置 workers2,将会生成两个线程并发加载数据。我们还应该将批次大小从默认的 1 增加到更合理的值,比如 64kBatchSize 的值)。那么让我们创建一个 DataLoaderOptions 对象并设置适当的属性:

auto data_loader = torch::data::make_data_loader(
    std::move(dataset),
    torch::data::DataLoaderOptions().batch_size(kBatchSize).workers(2));

我们现在可以编写一个循环来加载数据批次,并且暂时只将其打印到控制台:

for (torch::data::Example<>& batch : *data_loader) {
  std::cout << "Batch size: " << batch.data.size(0) << " | Labels: ";
  for (int64_t i = 0; i < batch.data.size(0); ++i) {
    std::cout << batch.target[i].item<int64_t>() << " ";
  }
  std::cout << std::endl;
}

The type returned by the data loader in this case is a torch::data::Example. This type is a simple struct with a data field for the data and a target field for the label. Because we applied the Stack collation earlier, the data loader returns only a single such example. If we had not applied the collation, the data loader would yield std::vector<torch::data::Example<>> instead, with one element per example in the batch.

如果你重新构建并运行这段代码,你应该会看到类似这样的结果:

root@fa350df05ecf:/home/build# make
Scanning dependencies of target dcgan
[ 50%] Building CXX object CMakeFiles/dcgan.dir/dcgan.cpp.o
[100%] Linking CXX executable dcgan
[100%] Built target dcgan
root@fa350df05ecf:/home/build# make
[100%] Built target dcgan
root@fa350df05ecf:/home/build# ./dcgan
Batch size: 64 | Labels: 5 2 6 7 2 1 6 7 0 1 6 2 3 6 9 1 8 4 0 6 5 3 3 0 4 6 6 6 4 0 8 6 0 6 9 2 4 0 2 8 6 3 3 2 9 2 0 1 4 2 3 4 8 2 9 9 3 5 8 0 0 7 9 9
Batch size: 64 | Labels: 2 2 4 7 1 2 8 8 6 9 0 2 2 9 3 6 1 3 8 0 4 4 8 8 8 9 2 6 4 7 1 5 0 9 7 5 4 3 5 4 1 2 8 0 7 1 9 6 1 6 5 3 4 4 1 2 3 2 3 5 0 1 6 2
Batch size: 64 | Labels: 4 5 4 2 1 4 8 3 8 3 6 1 5 4 3 6 2 2 5 1 3 1 5 0 8 2 1 5 3 2 4 4 5 9 7 2 8 9 2 0 6 7 4 3 8 3 5 8 8 3 0 5 8 0 8 7 8 5 5 6 1 7 8 0
Batch size: 64 | Labels: 3 3 7 1 4 1 6 1 0 3 6 4 0 2 5 4 0 4 2 8 1 9 6 5 1 6 3 2 8 9 2 3 8 7 4 5 9 6 0 8 3 0 0 6 4 8 2 5 4 1 8 3 7 8 0 0 8 9 6 7 2 1 4 7
Batch size: 64 | Labels: 3 0 5 5 9 8 3 9 8 9 5 9 5 0 4 1 2 7 7 2 0 0 5 4 8 7 7 6 1 0 7 9 3 0 6 3 2 6 2 7 6 3 3 4 0 5 8 8 9 1 9 2 1 9 4 4 9 2 4 6 2 9 4 0
Batch size: 64 | Labels: 9 6 7 5 3 5 9 0 8 6 6 7 8 2 1 9 8 8 1 1 8 2 0 7 1 4 1 6 7 5 1 7 7 4 0 3 2 9 0 6 6 3 4 4 8 1 2 8 6 9 2 0 3 1 2 8 5 6 4 8 5 8 6 2
Batch size: 64 | Labels: 9 3 0 3 6 5 1 8 6 0 1 9 9 1 6 1 7 7 4 4 4 7 8 8 6 7 8 2 6 0 4 6 8 2 5 3 9 8 4 0 9 9 3 7 0 5 8 2 4 5 6 2 8 2 5 3 7 1 9 1 8 2 2 7
Batch size: 64 | Labels: 9 1 9 2 7 2 6 0 8 6 8 7 7 4 8 6 1 1 6 8 5 7 9 1 3 2 0 5 1 7 3 1 6 1 0 8 6 0 8 1 0 5 4 9 3 8 5 8 4 8 0 1 2 6 2 4 2 7 7 3 7 4 5 3
Batch size: 64 | Labels: 8 8 3 1 8 6 4 2 9 5 8 0 2 8 6 6 7 0 9 8 3 8 7 1 6 6 2 7 7 4 5 5 2 1 7 9 5 4 9 1 0 3 1 9 3 9 8 8 5 3 7 5 3 6 8 9 4 2 0 1 2 5 4 7
Batch size: 64 | Labels: 9 2 7 0 8 4 4 2 7 5 0 0 6 2 0 5 9 5 9 8 8 9 3 5 7 5 4 7 3 0 5 7 6 5 7 1 6 2 8 7 6 3 2 6 5 6 1 2 7 7 0 0 5 9 0 0 9 1 7 8 3 2 9 4
Batch size: 64 | Labels: 7 6 5 7 7 5 2 2 4 9 9 4 8 7 4 8 9 4 5 7 1 2 6 9 8 5 1 2 3 6 7 8 1 1 3 9 8 7 9 5 0 8 5 1 8 7 2 6 5 1 2 0 9 7 4 0 9 0 4 6 0 0 8 6
...

这意味着我们成功地从MNIST数据集中加载了数据。

编写训练循环

让我们现在完成我们示例中的算法部分,并实现生成器和判别器之间的微妙互动。首先,我们将创建两个优化器,一个用于生成器,一个用于判别器。我们使用的优化器实现了Adam算法:

torch::optim::Adam generator_optimizer(
    generator->parameters(), torch::optim::AdamOptions(2e-4).betas(std::make_tuple(0.5, 0.5)));
torch::optim::Adam discriminator_optimizer(
    discriminator->parameters(), torch::optim::AdamOptions(5e-4).betas(std::make_tuple(0.5, 0.5)));

注意

截至本文撰写时,C++前端提供了实现Adagrad、Adam、LBFGS、RMSprop和SGD的优化器。 文档 中有最新的列表。

接下来,我们需要更新训练循环。我们将添加一个外部循环,确保每轮迭代中数据加载器都被用完,然后编写生成对抗网络(GAN)的训练代码:

for (int64_t epoch = 1; epoch <= kNumberOfEpochs; ++epoch) {
  int64_t batch_index = 0;
  for (torch::data::Example<>& batch : *data_loader) {
    // Train discriminator with real images.
    discriminator->zero_grad();
    torch::Tensor real_images = batch.data;
    torch::Tensor real_labels = torch::empty(batch.data.size(0)).uniform_(0.8, 1.0);
    torch::Tensor real_output = discriminator->forward(real_images).reshape(real_labels.sizes());
    torch::Tensor d_loss_real = torch::binary_cross_entropy(real_output, real_labels);
    d_loss_real.backward();

    // Train discriminator with fake images.
    torch::Tensor noise = torch::randn({batch.data.size(0), kNoiseSize, 1, 1});
    torch::Tensor fake_images = generator->forward(noise);
    torch::Tensor fake_labels = torch::zeros(batch.data.size(0));
    torch::Tensor fake_output = discriminator->forward(fake_images.detach()).reshape(fake_labels.sizes());
    torch::Tensor d_loss_fake = torch::binary_cross_entropy(fake_output, fake_labels);
    d_loss_fake.backward();

    torch::Tensor d_loss = d_loss_real + d_loss_fake;
    discriminator_optimizer.step();

    // Train generator.
    generator->zero_grad();
    fake_labels.fill_(1);
    fake_output = discriminator->forward(fake_images).reshape(fake_labels.sizes());
    torch::Tensor g_loss = torch::binary_cross_entropy(fake_output, fake_labels);
    g_loss.backward();
    generator_optimizer.step();

    std::printf(
        "\r[%2ld/%2ld][%3ld/%3ld] D_loss: %.4f | G_loss: %.4f",
        epoch,
        kNumberOfEpochs,
        ++batch_index,
        batches_per_epoch,
        d_loss.item<float>(),
        g_loss.item<float>());
  }
}

以上,我们首先在真实图像上评估鉴别器,对于真实图像,它应该分配一个较高的概率。为此,我们使用 torch::empty(batch.data.size(0)).uniform_(0.8, 1.0) 作为目标概率。

注意

我们选择在 0.8 和 1.0 之间均匀分布的随机值,而不是处处使用 1.0,以便使判别器训练更加稳健。这个技巧被称为标签平滑

在评估判别器之前,我们将其参数的梯度清零。计算损失后,通过调用d_loss.backward()反向传播损失来计算新的梯度。我们为假图像重复这个过程。而不是使用数据集中的图像,我们让生成器通过输入一批随机噪声来创建假图像。然后我们将这些假图像传递给判别器。这次我们希望判别器发出低概率值,最好是全部为零。一旦我们分别计算了真实图像和假图像的判别器损失,我们就可以通过一步来更新判别器优化器的参数。

要训练生成器,我们再次首先将其梯度清零,然后重新评估鉴别器在假图像上的表现。然而,这次我们希望鉴别器为这些假图像分配的概率非常接近一,这表明生成器能够生成足以欺骗鉴别器认为它们实际上是真实图像的图像(来自数据集)。为此,我们将 fake_labels 张量填充为全部为一。最后,我们通过优化器更新生成器的参数。

我们现在应该可以在CPU上训练我们的模型了。我们还没有任何代码来捕获状态或采样输出,但我们会马上添加这些代码。目前,让我们只是观察一下我们的模型在做些什么 – 我们稍后会根据生成的图像来验证这是否有意义。重新构建并运行应该会打印出类似的内容:

root@3c0711f20896:/home/build# make && ./dcgan
Scanning dependencies of target dcgan
[ 50%] Building CXX object CMakeFiles/dcgan.dir/dcgan.cpp.o
[100%] Linking CXX executable dcgan
[100%] Built target dcga
[ 1/10][100/938] D_loss: 0.6876 | G_loss: 4.1304
[ 1/10][200/938] D_loss: 0.3776 | G_loss: 4.3101
[ 1/10][300/938] D_loss: 0.3652 | G_loss: 4.6626
[ 1/10][400/938] D_loss: 0.8057 | G_loss: 2.2795
[ 1/10][500/938] D_loss: 0.3531 | G_loss: 4.4452
[ 1/10][600/938] D_loss: 0.3501 | G_loss: 5.0811
[ 1/10][700/938] D_loss: 0.3581 | G_loss: 4.5623
[ 1/10][800/938] D_loss: 0.6423 | G_loss: 1.7385
[ 1/10][900/938] D_loss: 0.3592 | G_loss: 4.7333
[ 2/10][100/938] D_loss: 0.4660 | G_loss: 2.5242
[ 2/10][200/938] D_loss: 0.6364 | G_loss: 2.0886
[ 2/10][300/938] D_loss: 0.3717 | G_loss: 3.8103
[ 2/10][400/938] D_loss: 1.0201 | G_loss: 1.3544
[ 2/10][500/938] D_loss: 0.4522 | G_loss: 2.6545
...

移动到GPU

虽然我们当前的脚本可以在CPU上正常运行,但大家都知道卷积在GPU上会快得多。让我们快速讨论一下如何将训练迁移到GPU上。为此,我们需要做两件事:将GPU设备规格传递给我们自己分配的张量,并通过C++前端所有张量和模块都具有的to()方法显式地将其他张量复制到GPU上。实现这两者的最简单方法是在训练脚本的顶层创建一个torch::Device实例,然后将该设备传递给像torch::zeros这样的张量工厂函数以及to()方法。我们可以从使用CPU设备开始:

// Place this somewhere at the top of your training script.
torch::Device device(torch::kCPU);

新的张量分配如

torch::Tensor fake_labels = torch::zeros(batch.data.size(0));

应将参数更新为将 device 作为最后一个参数:

torch::Tensor fake_labels = torch::zeros(batch.data.size(0), device);

对于来自MNIST数据集等我们无法创建的张量,我们必须插入显式的to()调用。这意味着

torch::Tensor real_images = batch.data;

变成

torch::Tensor real_images = batch.data.to(device);

并且我们的模型参数也应该移动到正确的设备上:

generator->to(device);
discriminator->to(device);

注意

如果张量已经在提供的设备to()上存在,那么该调用将是一个无操作。不会创建额外的副本。

到目前为止,我们只是让之前驻留在CPU的代码更加明确。 然而,现在也非常容易将其设备更改为CUDA设备:

torch::Device device(torch::kCUDA)

现在所有的张量都将驻留在GPU上,调用快速的CUDA内核进行所有操作,而我们无需更改任何下游代码。如果我们想要指定某个特定的设备索引,可以在Device构造函数的第二个参数中传递。如果我们想要不同的张量驻留在不同的设备上,可以分别传递设备实例(例如一个在CUDA设备0上,另一个在CUDA设备1上)。我们甚至可以动态地进行这种配置,这通常很有用,可以使我们的训练脚本更加便携:

torch::Device device = torch::kCPU;
if (torch::cuda::is_available()) {
  std::cout << "CUDA is available! Training on GPU." << std::endl;
  device = torch::kCUDA;
}

或者甚至

torch::Device device(torch::cuda::is_available() ? torch::kCUDA : torch::kCPU);

检查点和恢复训练状态

我们训练脚本中最后一个需要添加的增强功能是对模型参数的状态、优化器的状态以及一些生成的图像样本进行定期保存。如果在训练过程中计算机突然崩溃,前两者将允许我们恢复训练状态。对于长时间的训练会话,这一点至关重要。幸运的是,C++前端提供了API来序列化和反序列化模型、优化器状态以及单个张量。

这个核心API是torch::save(thing,filename)torch::load(thing,filename),其中thing可以是一个torch::nn::Module子类或我们在训练脚本中有的优化器实例对象Adam。让我们更新我们的训练循环,在一定间隔内保存模型和优化器的状态:

if (batch_index % kCheckpointEvery == 0) {
  // Checkpoint the model and optimizer state.
  torch::save(generator, "generator-checkpoint.pt");
  torch::save(generator_optimizer, "generator-optimizer-checkpoint.pt");
  torch::save(discriminator, "discriminator-checkpoint.pt");
  torch::save(discriminator_optimizer, "discriminator-optimizer-checkpoint.pt");
  // Sample the generator and save the images.
  torch::Tensor samples = generator->forward(torch::randn({8, kNoiseSize, 1, 1}, device));
  torch::save((samples + 1.0) / 2.0, torch::str("dcgan-sample-", checkpoint_counter, ".pt"));
  std::cout << "\n-> checkpoint " << ++checkpoint_counter << '\n';
}

where kCheckpointEvery 是一个整数,设置为类似 100 的值以每 100 个批次保存一次检查点,而 checkpoint_counter 是每次保存检查点时递增的计数器。

要恢复训练状态,您可以在创建所有模型和优化器之后但在训练循环之前添加如下几行代码:

torch::optim::Adam generator_optimizer(
    generator->parameters(), torch::optim::AdamOptions(2e-4).beta1(0.5));
torch::optim::Adam discriminator_optimizer(
    discriminator->parameters(), torch::optim::AdamOptions(2e-4).beta1(0.5));

if (kRestoreFromCheckpoint) {
  torch::load(generator, "generator-checkpoint.pt");
  torch::load(generator_optimizer, "generator-optimizer-checkpoint.pt");
  torch::load(discriminator, "discriminator-checkpoint.pt");
  torch::load(
      discriminator_optimizer, "discriminator-optimizer-checkpoint.pt");
}

int64_t checkpoint_counter = 0;
for (int64_t epoch = 1; epoch <= kNumberOfEpochs; ++epoch) {
  int64_t batch_index = 0;
  for (torch::data::Example<>& batch : *data_loader) {

检查生成的图像

我们的训练脚本现在已经完成。我们可以在 CPU 或 GPU 上开始训练我们的生成对抗网络 (GAN)。为了检查训练过程中的中间输出,我们在代码中添加了定期将图像样本保存到 "dcgan-sample-xxx.pt" 文件中的功能,我们可以编写一个小型的 Python 脚本来加载这些张量并使用 matplotlib 显示它们:

import argparse

import matplotlib.pyplot as plt
import torch


parser = argparse.ArgumentParser()
parser.add_argument("-i", "--sample-file", required=True)
parser.add_argument("-o", "--out-file", default="out.png")
parser.add_argument("-d", "--dimension", type=int, default=3)
options = parser.parse_args()

module = torch.jit.load(options.sample_file)
images = list(module.parameters())[0]

for index in range(options.dimension * options.dimension):
  image = images[index].detach().cpu().reshape(28, 28).mul(255).to(torch.uint8)
  array = image.numpy()
  axis = plt.subplot(options.dimension, options.dimension, 1 + index)
  plt.imshow(array, cmap="gray")
  axis.get_xaxis().set_visible(False)
  axis.get_yaxis().set_visible(False)

plt.savefig(options.out_file)
print("Saved ", options.out_file)

现在让我们训练我们的模型大约30个周期:

root@3c0711f20896:/home/build# make && ./dcgan                                                                                                                                10:17:57
Scanning dependencies of target dcgan
[ 50%] Building CXX object CMakeFiles/dcgan.dir/dcgan.cpp.o
[100%] Linking CXX executable dcgan
[100%] Built target dcgan
CUDA is available! Training on GPU.
[ 1/30][200/938] D_loss: 0.4953 | G_loss: 4.0195
-> checkpoint 1
[ 1/30][400/938] D_loss: 0.3610 | G_loss: 4.8148
-> checkpoint 2
[ 1/30][600/938] D_loss: 0.4072 | G_loss: 4.36760
-> checkpoint 3
[ 1/30][800/938] D_loss: 0.4444 | G_loss: 4.0250
-> checkpoint 4
[ 2/30][200/938] D_loss: 0.3761 | G_loss: 3.8790
-> checkpoint 5
[ 2/30][400/938] D_loss: 0.3977 | G_loss: 3.3315
...
-> checkpoint 120
[30/30][938/938] D_loss: 0.3610 | G_loss: 3.8084

并在图表中显示这些图像:

root@3c0711f20896:/home/build# python display.py -i dcgan-sample-100.pt
Saved out.png

这应该看起来类似于这样:

digits

数字!万岁!现在轮到你大显身手了:你能改进模型让这些数字看起来更好吗?

结论

本教程希望能为您提供一个易于理解的PyTorch C++前端概述。像PyTorch这样的机器学习库必然有一个非常广泛和详细的API。因此,有许多概念我们没有时间和空间在这里讨论。然而,我鼓励您尝试使用API,并在遇到困难时查阅我们的文档,特别是库API部分。另外,请记住,只要有可能,C++前端的设计和语义将遵循Python前端的设计和语义,您可以利用这一点来提高您的学习效率。

提示

您可以在此仓库中找到本教程中展示的完整源代码。

如您遇到任何问题或有任何疑问,您可以使用我们的 论坛GitHub 问题 来联系我们。

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源