今日应用
今日话题
【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+深度学习各方向资源整理。
文章来源
原文地址: 点我阅读全文
原文作者: 极市平台