扩展

在 Google Colab 中运行

Colab

下载笔记本

笔记本

在 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 函数。

扩展-PyTorch,前端-API

自定义函数教程：融合卷积和批归一化

学习如何创建一个自定义 autograd 函数，将批归一化融合到卷积中以提高内存使用效率。

扩展-PyTorch,前端-API

自定义 C++ 和 CUDA 扩展

使用 numpy 创建一个没有参数的神经网络层。然后使用 scipy 创建一个具有可学习权重的神经网络层。

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

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

调度器是 PyTorch 的一个内部组件，负责在您调用像 torch::add 这样的函数时，确定实际应该运行什么代码。

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

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

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

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

通过 PrivateUse1 促进新后端集成

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

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