此实现使用来自 PyTorch 的 nn 软件包来构建网络。 PyTorch autograd 使定义计算图和获取梯度变得容易,但是原始的 autograd 对于定义复杂的神经网络来说可能太低了。 这是 nn 软件包可以提供帮助的地方。 nn 包定义了一组模块,您可以将其视为神经网络层,该神经网络层从输入产生输出并且可能具有一些可训练的权重。
此实现使用来自 PyTorch 的 nn 软件包来构建网络。 PyTorch Autograd 使定义计算图和获取梯度变得容易,但是原始的 Autograd 对于定义复杂的神经网络来说可能太低了。 这是 nn 软件包可以提供帮助的地方。 nn 包定义了一组模块,您可以将其视为神经网络层,该神经网络层从输入产生输出并且可能具有一些可训练的权重。
至此,我们有了一个同时具有 CPU 和 CUDA 实现的操作员。 我们如何为它添加 autograd 支持? 您可能会猜到,我们将注册一个 autograd 内核(类似于[自定义 autograd 函数](cpp_autograd)教程中描述的内容)! 但是,有一个变数:与 CPU 和 CUDA 内核不同,autograd 内核需要*重新分发*:它需要回调分派器才能到达最终的 CPU 和 CUDA 实现。
至此,我们有了一个同时具有 CPU 和 CUDA 实现的操作员。 我们如何为它添加 Autograd 支持? 您可能会猜到,我们将注册一个 Autograd 内核(类似于[自定义 Autograd 函数](cpp_autograd)教程中描述的内容)! 但是,有一个变数:与 CPU 和 CUDA 内核不同,Autograd 内核需要*重新分发*:它需要回调分派器才能到达最终的 CPU 和 CUDA 实现。
因此,在编写 autograd 内核之前,让我们编写一个*调度函数*,该函数调用调度程序以为您的操作员找到合适的内核。 该函数构成了供您的操作员使用的公共 C++ API,实际上,PyTorch C++ API 中的所有张量函数都在后台完全以相同的方式调用了调度程序。 调度功能如下所示:
因此,在编写 Autograd 内核之前,让我们编写一个*调度函数*,该函数调用调度程序以为您的操作员找到合适的内核。 该函数构成了供您的操作员使用的公共 C++ API,实际上,PyTorch C++ API 中的所有张量函数都在后台完全以相同的方式调用了调度程序。 调度功能如下所示: