扩展

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

扩展

本节内容将深入探讨如何扩展 PyTorch 的功能。内容涵盖自定义操作、前端 API 以及 C++ 扩展和调度器使用等高级主题。

---

PyTorch 自定义算子主页

这是关于 PyTorch 中自定义运算符所有内容的着陆页。

扩展 PyTorch,前端 API,C++,CUDA

自定义 Python 算子

在 Python 中创建自定义运算符。适用于将 Python 函数黑盒化以供 torch.compile 使用。

扩展 PyTorch,前端 API,C++,CUDA

自定义 C++ 和 CUDA 算子

如何使用自定义 C++ 和 CUDA 运算符扩展 PyTorch。

扩展 PyTorch,前端 API,C++,CUDA

自定义函数教程：二阶导数

学习如何编写支持二阶导数的自定义 autograd Function。

扩展 PyTorch,前端 API

自定义函数教程：融合卷积和 Batch Norm

学习如何创建自定义 autograd Function，将 Batch Norm 融合到卷积中以提高内存使用率。

扩展 PyTorch,前端 API

在 C++ 中注册一个分发算子

调度器是 PyTorch 的内部组件，负责确定调用 torch::add 等函数时实际运行的代码。

扩展 PyTorch,前端 API,C++

在 C++ 中为新后端扩展调度器

学习如何扩展调度器以添加位于 pytorch/pytorch 仓库之外的新设备，并对其进行维护，以保持与原生 PyTorch 设备同步。

扩展 PyTorch,前端 API,C++

通过 PrivateUse1 促进新后端集成

学习如何集成一个位于 pytorch/pytorch 仓库之外的新后端，并对其进行维护，以保持与原生 PyTorch 后端同步。

扩展 PyTorch,前端 API,C++