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

AI最新资讯1个月前发布 tree
13 0 0

今日应用


今日话题


【PyTorch 奇淫技巧】如何在PyTorch中创建和使用Python自定义操作符
【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+深度学习各方向资源整理。

文章来源


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

© 版权声明

相关文章

暂无评论

暂无评论...