【PyTorch 奇淫技巧】如何在PyTorch中创建和使用Python自定义操作符

今日应用

淘声网

全球免费声音素材一站式获取，独创toSound“吐司”引擎，助你轻松找到游戏音效、影视配乐等丰富音频素材，免费授权个人和商业使用。

今日话题

【PyTorch 奇淫技巧】如何在PyTorch中创建和使用Python自定义操作符

 

重点标签 Python自定义操作符教程、PyTorch 2.4新特性、flash attention、图中断、自定义运算符、封装Python函数、训练支持、PIL库、crop功能、FakeTensor kernel、梯度公式、opcheck、可变运算符、numpy.sin

文章摘要

极市导读：本文为读者提供在PyTorch中创建和使用Python自定义操作符的教程，以及如何将PIL库的crop功能封装为自定义运算符的示例。

前言

在vllm项目中，通过使用`@torch.library.custom_op`装饰器，我们注意到了flash attention包的新特性，这是PyTorch 2.4版本之后引入的功能。本文将介绍如何利用官方教程来创建自定义操作符，并与PyTorch集成。

Python Custom Operators 教程

本教程将指导读者如何创建Python自定义运算符，并将其与PyTorch集成。这包括使用`torch.library.opcheck`进行测试。先决条件是安装PyTorch 2.4或更高版本。

创建自定义运算符的原因

– 将任意Python函数视为不透明的可调用对象，与`torch.compile`相对应。
– 为任意Python函数添加训练支持。

示例：封装PIL库的crop功能

本文将展示如何将PIL库的`crop`操作封装为自定义运算符。首先，定义`crop`和`display`函数，然后通过`torch.compile`测试其兼容性。为了在`torch.compile`中使用`crop`，需要将其封装为自定义运算符，并添加”FakeTensor kernel“。

自定义运算符的封装

使用`@torch.library.custom_op`定义新的自定义操作符，并使用`@crop.register_fake`添加FakeTensor kernel。这样，`crop`操作就可以在`torch.compile`中正常工作。

为crop添加训练支持

使用`torch.library.register_autograd`为运算符添加训练支持。通过定义`backward`和`setup_context`函数，并使用`register_autograd`指定梯度公式。

测试Python自定义运算符

使用`torch.library.opcheck`测试自定义运算符是否正确注册。这不会测试梯度的数学正确性，需要单独编写测试。

可变的Python自定义运算符

演示了如何将修改输入的Python函数封装为自定义运算符，例如使用`numpy.sin`函数。

总结

本文介绍了如何使用`torch.library.custom_op`创建与PyTorch子系统协同工作的Python自定义运算符。更多详细信息，请参阅相关文档和教程。

公众号后台回复“数据集”获取100+深度学习各方向资源整理。

文章来源

原文地址: 点我阅读全文
原文作者: 极市平台